Skip to main content

PyTorch Malware: Why torch.load is Unsafe

It's Just a Zip File with a Bomb Inside​

Many developers assume that a PyTorch model (.pt or .pth) is just a collection of tensors—numbers representing weights and biases.

In reality, the standard PyTorch save format is a Zip archive containing Pickle files. When you run torch.load(), you are implicitly running pickle.load(). This inherits all the Remote Code Execution (RCE) risks of Python's pickle module.

The "State Dict" Myth​

A common misconception is: "I'm only loading the state_dict, so I'm safe."

False. The state_dict itself is serialized using pickle. To read the dictionary of tensors, PyTorch must deserialize the file. If the file contains a malicious opcode sequence, the code executes before the weights are even assigned to your model.

Real-World Attacks​

We have seen attacks on Hugging Face where:

  1. An attacker clones a popular repository (e.g., Llama-2).
  2. They inject a payload into the pytorch_model.bin file.
  3. They upload it as a "Fine-tuned" version.
  4. Users download it, run torch.load, and their AWS credentials are exfiltrated.

Mitigation: weights_only=True​

Recent versions of PyTorch introduced the weights_only=True flag for torch.load(). This restricts the unpickler to only load basic types (tensors, strings, dicts).

However:

  1. It is not the default in many older codebases.
  2. It can break compatibility with complex custom models.
  3. It relies on an allowlist that might still have edge cases.

The Veritensor Approach​

Instead of relying on runtime checks, you should verify artifacts before they enter your environment.

Veritensor inspects the internal structure of PyTorch Zip archives. It recursively scans the embedded pickle files for unsafe imports.

# Signature detection example
- "CRITICAL: Unsafe import 'socket' detected in model weights"

By scanning your ~/.cache/huggingface folder, you can ensure no "sleeper" models are waiting to execute code on your GPU cluster.