From aca1bbc82ee60be161eb538990d6b8924e164d76 Mon Sep 17 00:00:00 2001 From: ruv Date: Tue, 3 Mar 2026 18:08:31 -0500 Subject: [PATCH] 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 --- references/wifi_densepose_pytorch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/references/wifi_densepose_pytorch.py b/references/wifi_densepose_pytorch.py index 4d3475c6..844bdc89 100644 --- a/references/wifi_densepose_pytorch.py +++ b/references/wifi_densepose_pytorch.py @@ -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'])