adding weights_only flag

MouseLand · Feb 7, 2025 · 8f53ebe · 8f53ebe
1 parent 16b675c
commit 8f53ebe
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/cellpose/resnet_torch.py b/cellpose/resnet_torch.py
@@ -268,11 +268,12 @@ def load_model(self, filename, device=None):
             device (torch.device, optional): The device to load the model on. Defaults to None.
         """
         if (device is not None) and (device.type != "cpu"):
-            state_dict = torch.load(filename, map_location=device)
+            state_dict = torch.load(filename, map_location=device, weights_only=True)
         else:
             self.__init__(self.nbase, self.nout, self.sz, self.mkldnn, self.conv_3D,
                           self.diam_mean)
-            state_dict = torch.load(filename, map_location=torch.device("cpu"))
+            state_dict = torch.load(filename, map_location=torch.device("cpu"), 
+                                    weights_only=True)
 
         if state_dict["output.2.weight"].shape[0] != self.nout:
             for name in self.state_dict():
@@ -318,7 +319,8 @@ def load_model(self, filename, device=None):
         else:
             self.__init__(self.nbase, self.nout, self.sz, self.mkldnn, self.conv_3D,
                           self.diam_mean)
-            state_dict = torch.load(filename, map_location=torch.device("cpu"), weights_only=True)
+            state_dict = torch.load(filename, map_location=torch.device("cpu"), 
+                                    weights_only=True)
 
         self.load_state_dict(state_dict)