vrr/Ruview
2
0
Fork 0
mirror of https://github.com/ruvnet/RuView.git synced 2026-04-26 13:10:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 13

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

2
references/wifi_densepose_pytorch.py
View file

@ -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'])