frank/ruvnet--RuView
1
0
Fork 0
mirror of https://github.com/ruvnet/RuView synced 2026-06-09 10:13:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: use weights_only=True for safe PyTorch model loading

Browse Source 
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 #106

Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruv
2026-03-03 18:08:31 -05:00
parent 2ad510782e
commit aca1bbc82e
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
references/wifi_densepose_pytorch.py
+1 -1
View File
@@ -441,7 +441,7 @@ class WiFiDensePoseTrainer:
}, path) }, path)
def load_model(self, 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.model.load_state_dict(checkpoint['model_state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])