suppress TypedStorage warning during nequip-package

Author: cw-tan

nequip/model/saved_models/package.py

@@ -51,15 +51,15 @@ def _get_package_metadata(imp) -> Dict[str, Any]:
 
 
 @contextlib.contextmanager
-def _suppress_package_importer_warnings():
+def _suppress_package_importer_exporter_warnings():
     # Ideally this ceases to exist or becomes a no-op in future versions of PyTorch
     with warnings.catch_warnings():
         # suppress torch.package TypedStorage warning
         warnings.filterwarnings(
             "ignore",
             message="TypedStorage is deprecated.*",
             category=UserWarning,
-            module="torch.package.package_importer",
+            module=r"torch\.package\.(package_exporter|package_importer)",
         )
         yield
 
@@ -104,7 +104,7 @@ def ModelFromPackage(package_path: str, compile_mode: str = _EAGER_MODEL_KEY):
 
     # === load model ===
     logger.info(f"Loading model from package file: {package_path} ...")
-    with _suppress_package_importer_warnings():
+    with _suppress_package_importer_exporter_warnings():
         # during `nequip-package`, we need to use the same importer for all the models for successful repackaging
         # see https://pytorch.org/docs/stable/package.html#re-export-an-imported-object
         if workflow_state == "package":
@@ -141,7 +141,7 @@ def ModelFromPackage(package_path: str, compile_mode: str = _EAGER_MODEL_KEY):
 
 def data_dict_from_package(package_path: str) -> AtomicDataDict.Type:
     """Load example data from a .nequip.zip package file."""
-    with _suppress_package_importer_warnings():
+    with _suppress_package_importer_exporter_warnings():
         imp = torch.package.PackageImporter(package_path)
         data = imp.load_pickle(package="model", resource="example_data.pkl")
     return data
@@ -159,7 +159,7 @@ def ModelTypeNamesFromPackage(package_path: str):
 
     _check_file_exists(file_path=package_path, file_type="package")
 
-    with _suppress_package_importer_warnings():
+    with _suppress_package_importer_exporter_warnings():
         imp = torch.package.PackageImporter(package_path)
         pkg_metadata = _get_package_metadata(imp)
 

nequip/scripts/package.py

@@ -5,7 +5,7 @@
 from nequip.model.saved_models.package import (
     _get_shared_importer,
     _get_package_metadata,
-    _suppress_package_importer_warnings,
+    _suppress_package_importer_exporter_warnings,
 )
 from nequip.model.saved_models import load_saved_model
 from nequip.model.utils import (
@@ -100,7 +100,7 @@ def main(args=None):
             "packed model file to inspect must end with the `.nequip.zip` extension"
         )
 
-        with _suppress_package_importer_warnings():
+        with _suppress_package_importer_exporter_warnings():
             imp = torch.package.PackageImporter(args.pkg_path)
             pkg_metadata = _get_package_metadata(imp)
 
@@ -250,7 +250,7 @@ def main(args=None):
             models_to_package.update({compile_mode: model})
 
         # == package ==
-        with _suppress_package_importer_warnings():
+        with _suppress_package_importer_exporter_warnings():
             with torch.package.PackageExporter(
                 args.output_path, importer=importers, debug=True
             ) as exp: