How to Save and Load PyTorch Model State Dict in Python

This code demonstrates how to save a PyTorch model's state dict to a file and load it back into a new model instance, verifying weights match.

Easy Python 3.8+ Aug 9, 2026 ML engineering pipelines 13 views 0 copies

pytorch state-dict model saving loading
Requires third-party packages — install first
pip install torch

Python code

32 lines
Python 3.8+ 
import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 8)
        self.fc2 = nn.Linear(8, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

if __name__ == "__main__":
    model = SimpleNet()
    state_dict = model.state_dict()
    
    # Save state dict to a temp file
    torch.save(state_dict, "model_weights.pth")
    
    # Create a new model and load saved state dict
    new_model = SimpleNet()
    loaded_state = torch.load("model_weights.pth")
    new_model.load_state_dict(loaded_state)
    
    # Verify weights match
    for key in state_dict:
        assert torch.equal(state_dict[key], new_model.state_dict()[key])
    
    print("State dict saved and loaded successfully.")
    print("Number of parameter tensors:", len(state_dict))
    print("Sample weight shape:", state_dict["fc1.weight"].shape)

Output

stdout 
State dict saved and loaded successfully.
Number of parameter tensors: 4
Sample weight shape: torch.Size([8, 4])

How it works

The state_dict method returns an OrderedDict mapping each trainable parameter name to its tensor. torch.save serializes this dict to disk using Python's pickle format, and torch.load reads it back. The load_state_dict method copies the tensor values into the model in-place, aligning them by parameter name. Asserting equality with torch.equal ensures exact value preservation. This pattern is the standard way to persist and restore trained model weights without shipping the full model class.

Common mistakes

  • Forgetting to call `model.eval()` before saving or loading when using dropout/batch norm
  • Calling `load_state_dict` without a matching model architecture, causing runtime errors
  • Using `torch.load` on untrusted files, which can execute arbitrary code
  • Saving the whole model instead of only the state dict, making the file less portable

Variations

  1. Use `torch.save(model.state_dict(), f)` with a file object like `io.BytesIO` to avoid disk writes
  2. Load with `weights_only=True` in PyTorch 2.0+ to improve security

Real-world use cases

  • Persisting trained model weights to disk for later inference or further training in a production ML pipeline.
  • Transferring learned weights from a pre-trained model to a new model instance with a different batch size or input shape.
  • Creating reproducible experiments by saving and comparing state dicts across training runs.

Sponsored

Run locally

This sample needs third-party packages, so it cannot run in the browser IDE. Copy the code above, install the packages shown at the top, then run it in your own Python environment.

More from ML engineering pipelines

Keep learning

Related tutorials and quizzes for this topic.

Tutorials

Quizzes