Skip to main content

PyTorch Malware: The Structural Flaws of torch.load

A dangerous misconception within the AI engineering community is that a PyTorch model (.pt, .pth, or pytorch_model.bin) is simply a flat binary file containing dense tensor matrices (weights and biases).

Architecturally, the modern PyTorch save format is a Zip archive containing multiple Pickle files. When an application executes torch.load(), it is implicitly extracting this archive and running pickle.load() on the internal data.pkl stream, inheriting all the Remote Code Execution (RCE) vulnerabilities inherent to the Pickle Virtual Machine.

The state_dict Safety Myth​

Developers frequently assume that loading only the weights is safe: "I am only loading the state_dict(), not the full model architecture, so I am protected against RCE."

This is mathematically false. The state_dict itself is a Python dictionary object serialized via Pickle. To map the tensor data into memory, the PyTorch engine must deserialize the file. If an adversary has injected a malicious __reduce__ opcode sequence into the archive, the PVM executes the payload before the dictionary is constructed and assigned to the model.

The weights_only=True Mitigation and Its Limits​

Recent iterations of PyTorch introduced a security parameter to mitigate this exact vector:

# Restricting the unpickler to safe data types
model_weights = torch.load("model.pt", weights_only=True)

Setting weights_only=True restricts the internal unpickler to a strict allowlist of fundamental data types (tensors, strings, integers, dicts), explicitly blocking the instantiation of arbitrary classes or system calls.

However, architectural blind spots remain:

  1. Legacy Codebases: Millions of inference scripts in production do not enforce this flag.

  2. Custom Architectures: Highly customized models frequently utilize complex Python objects that break when the strict allowlist is enforced.

  3. Bypass Vectors: Researchers continuously hunt for edge-case vulnerabilities within the allowlisted object graphs.

Structural Verification via Veritensor​

Instead of relying on runtime exception handling, enterprise security requires deterministic pre-execution validation.

Veritensor operates directly on the binary artifact. It natively parses the PyTorch Zip archive structure, extracts the embedded data.pkl, and statically emulates the PVM opcodes. By scanning your local .cache/huggingface directories or CI/CD model fetch pipelines, Veritensor identifies dangerous imports and structurally guarantees that the artifact is purely data-driven before it is ever loaded into VRAM.