Merge pull request #455 from mj-will/fix-torch-loading

MAINT: specify `weights_only=True` in `load.torch`
mj-will · Feb 3, 2025 · b1e88f3 · b1e88f3
2 parents 49ce6d8 + 04b3c80
commit b1e88f3
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/src/nessai/flowmodel/base.py b/src/nessai/flowmodel/base.py
@@ -726,7 +726,7 @@ def load_weights(self, weights_file):
         # TODO: these two methods are basically the same
         if not self.initialised:
             self.initialise()
-        self.model.load_state_dict(torch.load(weights_file))
+        self.model.load_state_dict(torch.load(weights_file, weights_only=True))
         self.model.eval()
         self.weights_file = weights_file
 

diff --git a/src/nessai/flowmodel/importance.py b/src/nessai/flowmodel/importance.py
@@ -160,7 +160,7 @@ def load_all_weights(self) -> None:
         for wf in self.weights_files:
             new_flow = configure_model(self.flow_config)
             new_flow.device = self.device
-            new_flow.load_state_dict(torch.load(wf))
+            new_flow.load_state_dict(torch.load(wf, weights_only=True))
             self.models.append(new_flow)
         self.models.eval()
 

diff --git a/tests/test_flowmodel/test_flowmodel_base.py b/tests/test_flowmodel/test_flowmodel_base.py
@@ -726,7 +726,7 @@ def test_load_weights(model, initialised):
         model.initialise.assert_not_called()
     else:
         model.initialise.assert_called_once()
-    mock_load.assert_called_once_with(weights_file)
+    mock_load.assert_called_once_with(weights_file, weights_only=True)
     model.model.load_state_dict.assert_called_once_with(d)
     model.model.eval.assert_called_once()
     assert model.weights_file == weights_file