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

fix: safe PyTorch model loading (weights_only=True)
2026-03-03 18:11:29 -05:00 · 2026-03-03 18:11:29 -05:00 · 26655d397e
parent 2ad510782e aca1bbc82e
commit 26655d397e
1 changed files with 1 additions and 1 deletions
--- a/references/wifi_densepose_pytorch.py
+++ b/references/wifi_densepose_pytorch.py
@ -441,7 +441,7 @@ class WiFiDensePoseTrainer:
        }, 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'])