Include Pytorch model files in the zip file filter logic

sassoftware · Jan 31, 2024 · eb4e267 · eb4e267
1 parent d4c013c
commit eb4e267
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/sasctl/pzmm/zip_model.py b/src/sasctl/pzmm/zip_model.py
@@ -33,6 +33,9 @@ def _filter_files(file_dir: Union[str, Path], is_viya4: Optional[bool] = False)
     file_names.extend(sorted(Path(file_dir).glob("*.pickle")))
     # Include H2O.ai MOJO files
     file_names.extend(sorted(Path(file_dir).glob("*.mojo")))
+    # Include Pytorch.ai files
+    file_names.extend(sorted(Path(file_dir).glob("*.pth")))
+    file_names.extend(sorted(Path(file_dir).glob("*.pt")))
     if file_names:
         return file_names
     else: