nn.Module: What Does PyTorch Think Belongs to Your Model?

Sat, 08 Aug 2026 13:09:00 +0100

Here is a classifier for the XOR-shaped dataset from Chapter 4, written the way most people write PyTorch. It trains. Validation accuracy climbs from chance to 95.7%. It saves a checkpoint, the checkpoint loads without a warning, and then the reloaded model performs at chance.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

n = 2000
X = torch.randn(n, 2)
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()

perm = torch.randperm(n)
X_train, y_train = X[perm[:1600]], y[perm[:1600]]
X_val, y_val = X[perm[1600:]], y[perm[1600:]]

class Encoder(nn.Module):
    def __init__(self, width=16, depth=3):
        super().__init__()
        self.blocks = [nn.Linear(width, width) for _ in range(depth)]

    def forward(self, x):
        for block in self.blocks:
            x = torch.relu(block(x))
        return x

class Classifier(nn.Module):
    def __init__(self, width=16, depth=3):
        super().__init__()
        self.stem = nn.Linear(2, width)
        self.encoder = Encoder(width, depth)
        self.head = nn.Linear(width, 2)

    def forward(self, x):
        x = torch.relu(self.stem(x))
        x = self.encoder(x)
        return self.head(x)

model = Classifier()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

for epoch in range(200):
    logits = model(X_train)
    loss = F.cross_entropy(logits, y_train)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 50 == 0 or epoch == 199:
        with torch.no_grad():
            val_acc = (model(X_val).argmax(1) == y_val).float().mean()
        print(f"epoch={epoch:3d} loss={loss.item():.4f} "
              f"val_acc={val_acc.item():.3f}")

torch.save(model.state_dict(), "classifier.pt")

reloaded = Classifier()
print(reloaded.load_state_dict(torch.load("classifier.pt", weights_only=True)))

with torch.no_grad():
    acc = (reloaded(X_val).argmax(1) == y_val).float().mean()
print(f"reloaded val_acc={acc.item():.3f}")

epoch=  0 loss=0.6964 val_acc=0.512
epoch= 50 loss=0.5930 val_acc=0.810
epoch=100 loss=0.3856 val_acc=0.873
epoch=150 loss=0.2550 val_acc=0.925
epoch=199 loss=0.1908 val_acc=0.957
<All keys matched successfully>
reloaded val_acc=0.512

Read the last two lines together. PyTorch reports that every key in the checkpoint matched every key the model expected, and the model that came back is worthless. Nothing raised. Nothing warned. load_state_dict was as satisfied as it is capable of being.

Registration on Programmer.ie

nn.Module: What Does PyTorch Think Belongs to Your Model?