frank/ruvnet--RuView
1
0
Fork 0
mirror of https://github.com/ruvnet/RuView synced 2026-08-07 20:01:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
fix: safe PyTorch model loading (weights_only=True)
This commit is contained in:
rUv
2026-03-03 18:11:29 -05:00
committed by GitHub
parent 2ad510782e aca1bbc82e
commit 26655d397e
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'])