fix: use weights_only=True for safe PyTorch model loading

ruvnet · ruvnet · commit aca1bbc82ee6 · 2026-03-03T18:08:31.000-05:00
Replace unsafe `torch.load(path)` with `torch.load(path, map_location=self.device, weights_only=True)` to prevent pickle deserialization RCE (trailofbits.python.pickles-in-pytorch). weights_only=True disables pickle entirely for model loading, which is the PyTorch-recommended mitigation (available since 1.13). Also adds map_location for correct CPU/GPU device mapping. Closes ruvnet#106 Co-Authored-By: claude-flow <ruv@ruv.net>
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'])
 
