Saving, Loading, and Deploying Models

Build the mental model

A trained model's entire learned knowledge lives in its parameters — the weights and biases inside each layer. PyTorch's recommended way to persist this is saving model.state_dict(), a plain dictionary mapping each layer's name to its parameter tensors, rather than pickling the whole model object. The naive approach, torch.save(model), ties the saved file to the exact class definition and file layout that existed at save time; if that code moves, changes, or is refactored later, loading can silently break or fail outright. Saving just the state_dict avoids this: you keep the model's architecture defined in code as a class you can always reconstruct, and loading becomes two explicit steps — instantiate a fresh copy of that architecture, then call load_state_dict() to copy the saved tensors into it. This makes checkpoints portable across code changes, more transparent to inspect, and safer to share. Equally important, and easy to forget, is that a freshly loaded model still needs model.eval() before you run predictions — it switches off training-only behavior like dropout (which randomly zeroes activations) and swaps batch normalization to its stored running statistics. Skip it, and inference becomes needlessly nondeterministic or biased.

Connect it to a real scenario

On the Tutorial Platform, the sentiment classifier trained for course reviews and the ranking model behind search shouldn't be retrained on every server restart or every request — they're trained once, then saved as checkpoints and loaded by the API server that answers requests. Saving state_dict() (not the whole object) means the model architecture can be refactored or moved between files later without breaking old checkpoints. And every inference-serving code path — the endpoint that scores a new review or ranks search results — must call .eval() right after loading; forgetting it on a model using dropout would make the same review get a different sentiment score on every request, exactly the kind of bug that's invisible in testing but obvious in production.

Save model.state_dict() (the learned parameters), not the whole model object — it's more portable and doesn't depend on the exact class definition being unpickled unchanged later.
Loading requires two steps: recreate the model architecture in code, then call load_state_dict() to copy the saved parameters into it.
Always call model.eval() before inference — it disables dropout and switches batch norm to stored running statistics, without which predictions can be nondeterministic or wrong.

Try the working example

python

import torch
import torch.nn as nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 8)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(8, 1)

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

torch.manual_seed(0)
model = TinyNet()

# Quick "training" so weights aren't just fresh init
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
x_train = torch.randn(5, 4)
y_train = torch.randn(5, 1)
for _ in range(20):
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x_train), y_train)
    loss.backward()
    optimizer.step()

# Save only the learned parameters, not the whole object
torch.save(model.state_dict(), "tiny_net.pt")

# Simulate a fresh process: recreate the architecture, then load weights
loaded_model = TinyNet()
loaded_model.load_state_dict(torch.load("tiny_net.pt"))
loaded_model.eval()  # disable dropout for deterministic inference

model.eval()
sample = torch.randn(1, 4)
with torch.no_grad():
    original_output = model(sample)
    loaded_output = loaded_model(sample)

print("Original model output:", original_output)
print("Loaded model output:  ", loaded_output)
print("Outputs match:", torch.allclose(original_output, loaded_output))

You should see

Original model output: tensor([[0.0421]])
Loaded model output:   tensor([[0.0421]])
Outputs match: True
(the exact numbers depend on the random seed and PyTorch version, but the two lines are always identical since both models share the same loaded weights and both run in eval mode)

5-minute try-it

Modify the code to add a second hidden layer with nn.BatchNorm1d after fc1, save/load it the same way, and verify that calling loaded_model.eval() (versus leaving it in train mode) changes the output — print both to see the difference dropout and batch norm make.

One important caution

Using torch.save(model) to pickle the whole object instead of state_dict — loading later can break if the class definition or file structure has changed.

Forgetting model.eval() after loading — leaving dropout active means the same input produces a different output on every call, and batch norm uses noisy per-batch statistics instead of the trained running averages.

PyTorch Docs — Saving and Loading Models — Deep Learning