Fixed PyTorch security warning by setting weights_only=True in torch.…

…load()
weigertlab · Aug 28, 2024 · d6d008b · d6d008b
1 parent 3b95f54
commit d6d008b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/spotiflow/model/spotiflow.py b/spotiflow/model/spotiflow.py
@@ -605,7 +605,7 @@ def load(
 
         states_path = Path(path) / f"{which}.pt"
 
-        checkpoint = torch.load(states_path, map_location=map_location)
+        checkpoint = torch.load(states_path, map_location=map_location, weights_only=True)
         model_state = self.cleanup_state_dict_keys(checkpoint["state_dict"])
         self.load_state_dict(model_state)