PyTorch Checkpoints with state_dict: Key Layout, Prefixes, and Strict Loading

The problem and the practical takeaway

You resume training or run inference with a saved PyTorch checkpoint and the script either raises a KeyError, reports missing or unexpected keys, or—worse—loads successfully while leaving some layers at random initialization. The cause is usually a mismatch between the checkpoint’s state_dict key layout and the live model’s key layout.

Takeaway: treat a checkpoint file as data, not as a drop-in replacement for the model class. Inspect the saved keys, normalize module. prefixes from DataParallel or DistributedDataParallel checkpoints, choose strict loading deliberately, and verify tensor shapes before trusting the restored weights.

Requirements

• The model class definition must be importable in the loading environment; state_dict stores tensor names and values, not the class code.
• PyTorch 2.x is assumed for the examples; torch.save and torch.load use Python’s pickle protocol by default.
• File-system read permission on the checkpoint and write permission on the destination directory when saving.
• Knowledge of whether the checkpoint was saved from a plain nn.Module, nn.DataParallel, or nn.parallel.DistributedDataParallel wrapper.

state_dict is an ordered mapping of parameter and buffer names to tensors. A buffer is a non-parameter tensor registered on a module, such as running statistics in batch normalization. strict controls whether load_state_dict raises on missing or unexpected keys.

Smallest suitable design

Save only model.state_dict(), not the whole model object. A whole-model torch.save(model, path) pickles the class reference and breaks more easily across PyTorch versions or when the class definition moves. Store the checkpoint alongside a version note recording PyTorch version, Python version, and the model class source location.

For training resumption, also save optimizer and scheduler state dictionaries in the same file as separate entries, but keep model weights under a clearly named key such as model.

Trust and data boundaries

torch.load executes arbitrary code during unpickling. Never load a checkpoint from an untrusted source without inspecting it first. Use weights_only=True when the file should contain only tensors and basic Python types; this blocks arbitrary object reconstruction.

When loading across devices, supply map_location. A checkpoint saved on CUDA does not automatically move to CPU, and a CPU checkpoint does not move to GPU.

Operational checks before loading

1. Open a Python REPL or script on the host that can read .
2. Load the checkpoint tensors only: torch.load("", weights_only=True).
3. Print the saved keys and the live model keys: print(list(saved.keys())) and print(list(model.state_dict().keys())).
4. Compare the two sets. Report missing keys (saved - live) and unexpected keys (live - saved).
5. For any shared key, compare shapes with saved[k].shape == model.state_dict()[k].shape. Use torch.equal only if you need to confirm identical values.
6. Load with strict=True first. Relax to strict=False only when you intentionally expect partial matches, and then audit which keys were skipped.

Risk: strict=False silently skips missing keys, so a model can load with untrained layers and produce plausible-looking but wrong outputs.

Failure modes

Failure	Cause	Check
KeyError or missing/unexpected keys	Architecture changed between save and load	Compare key sets before calling load_state_dict
Shape mismatch error	Same key name, different tensor dimensions	Compare .shape for each shared key
Device mismatch	CUDA checkpoint loaded on CPU or vice versa	Pass map_location to torch.load
Keys prefixed with module.	Checkpoint saved from DataParallel or DistributedDataParallel	Strip or add the prefix before loading
Silent partial load	strict=False hides missing keys	Audit the missing_keys and unexpected_keys return value
Pickle incompatibility	Major PyTorch version upgrade	Test round-trip save and load in the target environment

Conditions that change the design

• Moving from DistributedDataParallel to a single device: strip the module. prefix from saved keys before loading into an unwrapped model.
• Moving from a single device to DistributedDataParallel: add the module. prefix, or wrap the model before loading.
• Inference-only deployment: drop optimizer and scheduler state; keep only model weights and, if needed, class-to-index mappings.
• Custom attributes that are not buffers or parameters: state_dict does not capture them; implement custom state_dict/load_state_dict or save a sidecar JSON file.
• Major PyTorch upgrade: re-serialize checkpoints in the new version and verify a forward pass produces identical outputs in evaluation mode before retiring old files.

Example: save, inspect, and load with prefix normalization

Run this in a Python script or REPL on the machine owning the checkpoint file. Replace with your actual path.

import torch
import torch.nn as nn


class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(4, 2)

    def forward(self, x):
        return self.fc(x)


# 1. Save only the state dictionary.
model = TinyNet()
torch.save(model.state_dict(), "")

# 2. Load tensors only and inspect keys before touching the model.
saved = torch.load("", weights_only=True)
saved_keys = list(saved.keys())
live_keys = list(TinyNet().state_dict().keys())
print("saved:", saved_keys)
print("live :", live_keys)

# 3. Load into a matching architecture with strict checking.
fresh = TinyNet()
fresh.load_state_dict(saved, strict=True)

# 4. Normalize a DataParallel / DistributedDataParallel checkpoint.
#    Keys look like "module.fc.weight" instead of "fc.weight".
ddp_saved = torch.load("", weights_only=True)
cleaned = {k.replace("module.", "", 1): v for k, v in ddp_saved.items()}
fresh.load_state_dict(cleaned, strict=True)

# 5. Practical verification: deterministic forward pass before and after.
fresh.eval()
x = torch.randn(1, 4)
with torch.no_grad():
    out_before = fresh(x)
# Re-load into a newly constructed model and run the same input;
# outputs should match when weights and key mapping are correct.

Expected check: the printed saved and live key lists should match exactly for a plain model. For a DDP checkpoint, every saved key should begin with module. before stripping, and every cleaned key should exist in the live model.

Limitations and how to verify the result

state_dict does not capture optimizer momentum, learning-rate schedules, random-number-generator state, or custom Python attributes attached to a module. It also does not store the model class; if the class definition is missing or renamed, loading fails or maps weights to the wrong architecture.

Practical verification: after loading, set the model to eval(), run a fixed input through it, and save the output. Reload the checkpoint into a newly constructed model and run the same input. The outputs should match in deterministic mode. If they differ, the key mapping or strictness setting is wrong.

When saving, copy any existing checkpoint file to a backup name before overwriting it, because torch.save truncates the destination file and there is no built-in rollback for a failed write.