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

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
+1 -1
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'])