Loading PyTorch Models

The practical counterpart to Saving PyTorch Models — correctly loading saved weights and checkpoints back, including the architecture-mismatch pitfalls that most commonly trip up beginners.

Loading Weights Into a Matching Architecture

import torch

model = MyModelClass()   # the architecture MUST match the saved weights exactly
model.load_state_dict(torch.load("model_weights.pt"))
model.eval()               # switch to evaluation mode before inference

Resuming From a Full Checkpoint

checkpoint = torch.load("checkpoint_epoch_10.pt")

model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
start_epoch = checkpoint['epoch'] + 1

for epoch in range(start_epoch, num_epochs):
    # training continues exactly where it left off
    pass

Loading to a Different Device Than It Was Saved From

# A model saved from a GPU machine, loaded on a CPU-only machine
model.load_state_dict(torch.load("model_weights.pt", map_location=torch.device('cpu')))

Without map_location, loading GPU-saved weights on a machine without a GPU raises an error — this argument explicitly tells PyTorch which device to place the loaded tensors on, regardless of where they were originally saved from.

Loading Partial Weights (for Transfer Learning)

pretrained_dict = torch.load("pretrained_weights.pt")
model_dict = model.state_dict()

matched_dict = {k: v for k, v in pretrained_dict.items()
                 if k in model_dict and v.shape == model_dict[k].shape}
model_dict.update(matched_dict)
model.load_state_dict(model_dict)   # loads whatever matches; leaves the rest at its current initialization

This is the exact pattern from Model Saving and Loading, useful when loading a pretrained backbone into a model with a newly added, differently-shaped output layer.

Common Mistakes

• Forgetting model.eval() after loading for inference — dropout and batch normalization will still behave in training mode.
• Attempting to load a state_dict into a model with a mismatched architecture without using the partial-loading pattern — this raises a key/shape mismatch error rather than silently doing something wrong, which is at least easy to catch, but confusing the first time.
• Forgetting map_location when loading GPU-saved weights on a CPU-only environment.

Interview Relevance

Q: "What happens if you try to load a state_dict saved from a model with 3 hidden layers into a freshly created model with only 2 hidden layers?" load_state_dict raises an error about missing or unexpected keys, since it tries to match every parameter name in the saved file against the current model's parameter names exactly — a structural mismatch like a different number of layers means some saved keys won't have a corresponding destination (or vice versa). Loading would need the partial-loading pattern (filtering to only matching keys/shapes) to succeed with a genuinely different architecture.

Practice Question

Why is map_location needed when loading a model checkpoint saved on a GPU machine onto a machine with no GPU available?

Related DL Notes

• Saving PyTorch Models
• PyTorch Transfer Learning
• Model Saving and Loading