adding weights_only flag

Author: carsen-stringer

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)