From 5dbd768cb0c9f59ba96148007347ab0acb188a9e Mon Sep 17 00:00:00 2001 From: GEAK RDNA Smoke Date: Wed, 26 Aug 2026 06:49:29 -0400 Subject: [PATCH] Default to weights_only=True when loading pretrained weights WeightsEnum.get_state_dict forwards to torch.hub.load_state_dict_from_url, which passes weights_only explicitly to torch.load (defaulting to False). Because it is passed explicitly, this path does not pick up the weights_only=True default that torch.load adopted in PyTorch 2.6, so a tampered checkpoint can still execute arbitrary code via pickle during unpickling on any torch version. Set weights_only=True by default via setdefault so every pretrained checkpoint is loaded safely while leaving callers free to override. Verified that float and quantized (fbgemm) checkpoints still load, since the tensor/qtensor rebuild functions they use are on torch's weights_only allowlist. Co-Authored-By: Claude Opus 5 (1M context) --- torchvision/models/_api.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/torchvision/models/_api.py b/torchvision/models/_api.py index 5f46bcb5cd0..ff7f6cb967a 100644 --- a/torchvision/models/_api.py +++ b/torchvision/models/_api.py @@ -88,6 +88,13 @@ def verify(cls, obj: Any) -> Any: return obj def get_state_dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + # Load checkpoints with ``weights_only=True`` by default so a tampered + # checkpoint cannot execute arbitrary code via pickle during unpickling. + # ``torch.hub.load_state_dict_from_url`` forwards ``weights_only`` to + # ``torch.load`` explicitly, so it does not benefit from the + # ``weights_only=True`` default that ``torch.load`` itself adopted in + # PyTorch 2.6 -- it must be requested here. Callers may still override. + kwargs.setdefault("weights_only", True) return load_state_dict_from_url(self.url, *args, **kwargs) def __repr__(self) -> str: