Why torch.load can run code, and safetensors

Setup#

pip install torch safetensors

Everything runs on CPU. The demonstration below is deliberately harmless — it prints a line where a real attack would do something you would not enjoy. The mechanism is identical.

Watching a checkpoint execute code#

import re
import torch

class Sneaky:
    """A pickle stores instructions, not data. This one calls print; a real
    attack calls something worse. The mechanism is identical."""
    def __reduce__(self):
        return (print, ("*** code from the checkpoint file just ran ***",))

torch.save({"weights": torch.zeros(2), "note": Sneaky()}, "sneaky.pt")

print("loading with weights_only=False:")
torch.load("sneaky.pt", weights_only=False)

print("\nloading with weights_only=True:")
try:
    torch.load("sneaky.pt", weights_only=True)
except Exception as e:
    clean = re.sub(r"\x1b\[[0-9;]*m", "", str(e))      # strip terminal colours
    print(type(e).__name__)
    print(clean.splitlines()[0])
    print([l for l in clean.splitlines() if "Unsupported global" in l][0].strip())

loading with weights_only=False:
*** code from the checkpoint file just ran ***

loading with weights_only=True:
UnpicklingError
Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
WeightsUnpickler error: Unsupported global: GLOBAL builtins.print was not an allowed global by default. Please use `torch.serialization.add_safe_globals([print])` to allowlist this global if you trust this class/function.

Notice what the first load did not do. It did not warn, it did not slow down, and it produced a perfectly valid dictionary with the right weights in it. The code ran as a side effect of opening the file.

__reduce__ is a standard, documented part of Python's pickle protocol. It returns a callable and its arguments, and the unpickler calls it. Nothing here is a bug or an exploit of a flaw — it is the format working as designed.

The weights_only unpickler has an allowlist of types it will construct: tensors, storages, OrderedDict, and primitives. Anything else raises, naming the exact global it refused. That error message is what a blocked attack looks like.

Reading a file before you trust it#

import pickletools
import zipfile

with zipfile.ZipFile("sneaky.pt") as z:                 # a .pt file is a zip
    members = z.namelist()
    name = [n for n in members if n.endswith("data.pkl")][0]
    ops = [op.name for op, _, _ in pickletools.genops(z.read(name))]

print("archive members:", len(members))
print("GLOBAL/STACK_GLOBAL opcodes:", ops.count("GLOBAL") + ops.count("STACK_GLOBAL"))
print("REDUCE opcodes:", ops.count("REDUCE"))

archive members: 5
GLOBAL/STACK_GLOBAL opcodes: 4
REDUCE opcodes: 3

GLOBAL opcodes name a function to import; REDUCE calls one. Their presence is not proof of anything — normal checkpoints contain several, since rebuilding a tensor is itself a function call. But the names they import are informative, and a file importing os or subprocess deserves a hard look. This is roughly what a model hub's scanner does at scale.

safetensors: the format with nowhere to hide#

import torch
import torch.nn as nn
from safetensors.torch import save_file, load_file, safe_open

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
save_file(model.state_dict(), "model.safetensors")

# read the header without loading a single weight
with safe_open("model.safetensors", framework="pt") as f:
    print("keys:", f.keys())
    print("shape of 0.weight:", f.get_slice("0.weight").get_shape())
    print("one row only:", f.get_slice("0.weight")[0:1].shape)

loaded = load_file("model.safetensors")
fresh = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
print("load report:", fresh.load_state_dict(loaded))

keys: ['0.bias', '0.weight', '2.bias', '2.weight']
shape of 0.weight: [8, 4]
one row only: torch.Size([1, 4])
load report: <All keys matched successfully>

Three properties fall out of the format.

No code. The file is a JSON header plus raw bytes. There is no opcode, no callable, nothing to execute.

Lazy. safe_open read the shapes without loading the weights, and get_slice(...)[0:1] pulled one row off the disk. That is how a sharded 70-billion-parameter model gets loaded piece by piece without needing the whole thing in RAM.

Fast. The tensors are contiguous and aligned, so loading can memory-map rather than deserialise. On large models this is the difference between minutes and seconds.

What it does not do: safetensors stores tensors and a small string-to-string metadata map, and nothing else. Your optimizer state, epoch counter and config dictionary do not fit. Production setups commonly use safetensors for the weights people download and an ordinary .pt for the training state nobody else opens.

A policy that is short enough to follow#

# your own checkpoints, written by your own job
torch.save(state, "run.pt")
state = torch.load("run.pt", weights_only=True, map_location="cpu")

# anything downloaded
from safetensors.torch import load_file
model.load_state_dict(load_file("downloaded.safetensors"))

If a downloaded model exists only as .bin or .pt, and weights_only=True refuses it, the right response is to ask why the file needs to construct a custom class — not to flip the flag off. torch.serialization.add_safe_globals([...]) lets you allow specific known classes without opening the door to everything.

Common mistakes#

Passing weights_only=False to make an error go away. That error is the control working. Find out what the file wants to construct first.

Believing a scanner clears a file. Scanners catch known patterns. They are a filter, not a proof.

Trusting a file because the weights are correct. Both things happen. Correct weights are not evidence of an absent payload.

Assuming safetensors is slower because it is safer. It is generally faster to load, because there is less work to do.

Running downloaded model code because the weights were safe. A safetensors file plus a modeling_custom.py that you import is back to running someone else's code. The file format protects the weights, not your import statements — the same reasoning as prompt injection: the safe channel does not make the surrounding system safe.

Try it yourself#

Change Sneaky.__reduce__ to return (len, ("hello",)) and print what torch.load(..., weights_only=False) gives back for the "note" key. Then run pickletools.dis on the pickle and find the REDUCE opcode that did it.

What to learn next#

• Exponential moving average of weights — a second set of weights worth shipping.
• TorchScript: tracing vs scripting — packaging a model so the class definition is no longer needed.
• Model cards — documenting provenance so the next person can judge the file.