Skip to content

Commit

Permalink
Merge pull request #14 from ajinkya-kulkarni/patch-2
Browse files Browse the repository at this point in the history
Fixed PyTorch security warning by setting weights_only=True in torch.…
  • Loading branch information
maweigert authored Aug 28, 2024
2 parents 3b95f54 + d6d008b commit 54f79c1
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion spotiflow/model/spotiflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit 54f79c1

Please sign in to comment.