CVE-2026-31214
CriticalSummary
The torch-checkpoint-shrink.py script in the ml-engineering project contains an insecure deserialization vulnerability. The use of torch.load() to process PyTorch checkpoint files (.pt) without enabling the weights_only=True parameter allows for the deserialization of arbitrary Python objects.
Risk Assessment
An attacker can exploit this vulnerability by providing a maliciously crafted checkpoint file, leading to arbitrary code execution in the context of the user running the script.
Recommendation
It is recommended to enable the weights_only=True parameter in the torch.load() function and to review and validate checkpoint files before use.
Original NVD description (English source)
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.