Learn AI Create the Future — with Pranay
Main site

Error database

Error(s) in loading state_dict: Missing key(s) / Unexpected key(s)

The checkpoint's layer names or shapes do not match the model you built. Look at the key prefixes first — a "module." prefix means DataParallel saved it, and a nested dict means you saved more than the state_dict.

The message you saw
Error(s) in loading state_dict: Missing key(s) / Unexpected key(s)

By Updated

The error

Output
RuntimeError: Error(s) in loading state_dict for ResNet:
    Missing key(s) in state_dict: "conv1.weight", "bn1.weight", "bn1.bias"
    Unexpected key(s) in state_dict: "module.conv1.weight", "module.bn1.weight", "module.bn1.bias"

Shape conflicts produce the third variant:

Output
    size mismatch for fc.weight: copying a param with shape torch.Size([1000, 512]) from checkpoint, the shape in current model is torch.Size([10, 512]).

What it means

A state_dict is a dictionary mapping layer names to weight tensors. load_state_dict matches names exactly: every layer in your model must find its tensor, with the right shape. "Missing" keys are layers your model has but the file lacks; "unexpected" keys are tensors in the file with no home in your model. The two lists, read together, tell you precisely how the model and checkpoint disagree.

Why it happens

Read the keys before anything else — the pattern identifies the cause:

  • Every unexpected key starts with module. → the checkpoint was saved from a model wrapped in DataParallel/DistributedDataParallel, which adds that prefix.
  • Unexpected keys like optimizer, epoch, model_state_dict → you saved a training checkpoint dict, and passed the whole dict instead of the model part.
  • Missing and unexpected keys with genuinely different names → the model class you built is not the architecture that was saved (edited layers, renamed attributes, different version of the model code).
  • Size mismatch on the final layer only → same architecture, different number of output classes.

How to fix it

1. For the module. prefix, strip it.

python
sd = torch.load("model.pt", map_location="cpu", weights_only=True)
sd = {k.removeprefix("module."): v for k, v in sd.items()}
model.load_state_dict(sd)

Or avoid the prefix at save time: torch.save(model.module.state_dict(), path) when the model is wrapped.

2. For a nested checkpoint dict, load the right entry.

python
ckpt = torch.load("checkpoint.pt", map_location="cpu", weights_only=True)
print(ckpt.keys())                        # e.g. dict_keys(['model_state_dict', 'optimizer', 'epoch'])
model.load_state_dict(ckpt["model_state_dict"])

3. For genuinely different names, rebuild the exact architecture. The class definition must match what was saved — same layers, same attribute names, same sizes. Print a few keys from each side and diff them:

python
print(list(sd.keys())[:5])
print(list(model.state_dict().keys())[:5])

4. For a changed final layer (fine-tuning), load everything else.

python
sd.pop("fc.weight", None)
sd.pop("fc.bias", None)
model.load_state_dict(sd, strict=False)

Use strict=False only with a curated state_dict like this. Blanket strict=False on a mismatched checkpoint loads garbage silently — the model runs and predicts nonsense.

How to prevent it

Save plain state_dicts (or a dict of them), never whole model objects. Keep the model-building code in one function used by both training and loading. After any load, run one known input through the model and check the output looks sane — a five-line smoke test catches silent mismatch.