Skip to content

fix: add weights_only=True to torch.load calls (CVE-2026-31253)#2585

Open
factnn wants to merge 1 commit into
Dao-AILab:mainfrom
factnn:fix/unsafe-torch-load
Open

fix: add weights_only=True to torch.load calls (CVE-2026-31253)#2585
factnn wants to merge 1 commit into
Dao-AILab:mainfrom
factnn:fix/unsafe-torch-load

Conversation

@factnn
Copy link
Copy Markdown

@factnn factnn commented May 23, 2026

Summary

torch.load without weights_only=True allows arbitrary code execution via crafted pickle payloads. This adds weights_only=True to all four torch.load call sites flagged by CVE-2026-31253:

  • training/src/utils/checkpoint.py
  • training/src/eval.py
  • flash_attn/utils/pretrained.py
  • flash_attn/models/llama.py

All four sites load model state dicts (nested dict[str, Tensor] structures), which are fully supported by the safe unpickler.

Fixes #2583

…load without weights_only=True allows arbitrary code execution via crafted pickle payloads. Add weights_only=True to all torch.load call sites to mitigate this risk.Fixes Dao-AILab#2583
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.

Unsafe torch.load in FlashAttention training / eval checkpoint paths

1 participant