fix: use weights_only=True for safe PyTorch model loading
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:
parent
2ad510782e
commit
aca1bbc82e
|
|
@ -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'])
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue