Merge pull request ruvnet#133 from ruvnet/fix/pickle-deserialization-safety

ruvnet · web-flow · commit 26655d397e6b · 2026-03-03T18:11:29.000-05:00
fix: safe PyTorch model loading (weights_only=True)
diff --git a/references/wifi_densepose_pytorch.py b/references/wifi_densepose_pytorch.py
@@ -441,7 +441,7 @@ def save_model(self, path):
         }, path)
     
     def load_model(self, path):
-        checkpoint = torch.load(path)
+        checkpoint = torch.load(path, map_location=self.device, weights_only=True)
         self.model.load_state_dict(checkpoint['model_state_dict'])
         self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
 
