From 928d8b6ba30616d31262bd2e4e81dd5c3d161536 Mon Sep 17 00:00:00 2001 From: cyberkittens Date: Tue, 28 Apr 2026 14:02:04 +0300 Subject: [PATCH] fix: prevent arbitrary code execution via torch.load Add weights_only=True to torch.load() call. Without this flag, loading a malicious checkpoint executes arbitrary Python via pickle deserialization (CWE-502). Reported via Google AIVRP issue 500737486. --- scripts/run_multimodal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/run_multimodal.py b/scripts/run_multimodal.py index 231e340..1ed3555 100644 --- a/scripts/run_multimodal.py +++ b/scripts/run_multimodal.py @@ -113,7 +113,7 @@ def main(_): device = torch.device(_DEVICE.value) with _set_default_tensor_type(model_config.get_dtype()): model = gemma3_model.Gemma3ForMultimodalLM(model_config) - model.load_state_dict(torch.load(_CKPT.value)['model_state_dict']) + model.load_state_dict(torch.load(_CKPT.value, weights_only=True)['model_state_dict']) # model.load_weights(_CKPT.value) model = model.to(device).eval() print('Model loading done')