Unsafe loading

stefdoerr · stefdoerr · commit f5d67d6ac7f0 · 2024-12-03T11:49:21.000+02:00
diff --git a/torchmdnet/models/model.py b/torchmdnet/models/model.py
@@ -209,7 +209,7 @@ def load_model(filepath, args=None, device="cpu", return_std=False, **kwargs):
             filepath, args=args, device=device, return_std=return_std, **kwargs
         )
     assert isinstance(filepath, str)
-    ckpt = torch.load(filepath, map_location="cpu", weights_only=True)
+    ckpt = torch.load(filepath, map_location="cpu", weights_only=False)
     if args is None:
         args = ckpt["hyper_parameters"]
 

Original file line number	Diff line number	Diff line change
`@@ -209,7 +209,7 @@ def load_model(filepath, args=None, device="cpu", return_std=False, **kwargs):`
`209`	`209`	`filepath, args=args, device=device, return_std=return_std, **kwargs`
`210`	`210`	`)`
`211`	`211`	`assert isinstance(filepath, str)`
`212`		`- ckpt = torch.load(filepath, map_location="cpu", weights_only=True)`
	`212`	`+ ckpt = torch.load(filepath, map_location="cpu", weights_only=False)`
`213`	`213`	`if args is None:`
`214`	`214`	`args = ckpt["hyper_parameters"]`
`215`	`215`