Skip to content

fix: add weights_only=True to torch.load in run_multimodal.py#103

Open
cyberkittens4u wants to merge 1 commit into
google:mainfrom
cyberkittens4u:fix/torch-load-weights-only
Open

fix: add weights_only=True to torch.load in run_multimodal.py#103
cyberkittens4u wants to merge 1 commit into
google:mainfrom
cyberkittens4u:fix/torch-load-weights-only

Conversation

@cyberkittens4u
Copy link
Copy Markdown

Summary

scripts/run_multimodal.py line 116 calls torch.load() without weights_only=True. PyTorch uses pickle for checkpoint deserialization by default; a maliciously crafted checkpoint file can execute arbitrary code when loaded, giving an attacker full RCE on any machine running this script with a controlled checkpoint path.

Vulnerable call:

model.load_state_dict(torch.load(_CKPT.value)['model_state_dict'])

Fix:

model.load_state_dict(torch.load(_CKPT.value, weights_only=True)['model_state_dict'])

weights_only=True restricts pickle to safe tensor types only, eliminating the deserialization attack surface.

This issue was reported through Google's AI Vulnerability Rewards Program (AIVRP), issue 500737486.

Sibling scripts run.py and run_xla.py are not affected — they use model.load_weights() which has a safe code path.

Refs: PyTorch security advisory on torch.load, CWE-502 (Deserialization of Untrusted Data).

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.
@google-cla
Copy link
Copy Markdown

google-cla Bot commented Apr 28, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant