MITRE service request: 1987825
Status: RESERVED (pending a qualifying public reference per CNA Rules §5.3).
The torch-checkpoint-shrink.py script in the ml-engineering project in commit 0099885db36a8f06556efe1faf552518852cb1e0 (2025-20-27) contains an insecure deserialization vulnerability (CWE-502). The script uses torch.load() to process PyTorch checkpoint files (.pt) without enabling the security-restrictive weights_only=True parameter. This oversight allows the deserialization of arbitrary Python objects via the pickle module. A remote attacker can exploit this by providing a maliciously crafted checkpoint file, leading to arbitrary code execution in the context of the user running the script.
The checkpoint maintenance utility loads every matching .pt file under a user-supplied directory using torch.load(..., map_location=device) with the unsafe default (weights_only=False). A malicious checkpoint can embed pickle gadgets (for example via __reduce__) so that merely shrinking or re-saving models executes attacker-controlled Python on the machine of the operator.
training/checkpoints/torch-checkpoint-shrink.py (checkpoint shrink workflow).0099885db36a8f06556efe1faf552518852cb1e0.torch-checkpoint-shrink.py with -checkpoint_dir and -patterns (for example .pt). The main module resolves arguments, calls checkpoint_shrink(), iterates matching files, and for each file calls shrink_pt_file().shrink_pt_file() loads each checkpoint with torch.load(f, map_location=device) without weights_only=True, so full pickle-style object graphs may run during load..pt / .ckpt from an untrusted repo, shared drive, or message, then runs the shrink script against that directory. No server exposure is required; impact is local code execution as the user running Python./tmp or running shell commands), consistent with unsafe PyTorch unpickling.High for any workflow that points this script at untrusted checkpoints. Impact is arbitrary code execution with the privileges of the interactive user or batch job account.
weights_only=True where the file format allows, or avoid torch.load on untrusted paths entirely.