Test, predict & save

The loss falling tells us the network learned the training data. But the real question is whether it learned to read digits in general — including ones it has never seen. That's what the 10,000 held-out test images from Chapter 3 are for.

Grade it on unseen digits

Add an evaluation pass to the end of train.py. Two new things here: model.eval() switches the model into inference mode, and torch.no_grad() tells PyTorch to skip gradient tracking — we're only measuring now, not learning, so this is faster:

model.eval()
correct = 0
with torch.no_grad():
  for images, labels in test_loader:
      images, labels = images.to(device), labels.to(device)
      preds = model(images)
      correct += (preds.argmax(1) == labels).sum().item()

accuracy = correct / len(test_data)
print(f"test accuracy: {accuracy:.2%}")

preds.argmax(1) picks the highest-scoring digit for each image; we count how many match the true label. Run python train.py again:

test accuracy: 97.31%

~97% — the network reads handwriting it has never seen, correctly, 97 times out of 100. From random noise, in under a minute, on your own machine. That's the whole game of supervised learning in miniature.

Predict a single image

Accuracy over 10,000 images is abstract. Let's watch it read one digit:

image, true_label = test_data[0]
with torch.no_grad():
  scores = model(image.unsqueeze(0).to(device))  # add a batch dimension
  guess = scores.argmax(1).item()

print(f"model guessed: {guess} · actual: {true_label}")

unsqueeze(0) turns the single 1×28×28 image into a batch of one (1×1×28×28), because the model always expects a batch. You should see it guess the right digit.

Save your work

Training gave you a set of learned weights. Save them so you never have to retrain — this writes a small file you can reload in any future script:

torch.save(model.state_dict(), "digits.pt")
print("saved to digits.pt")

Loading them back later is two lines — build the same architecture, then pour the weights in:

model = DigitClassifier().to(device)
model.load_state_dict(torch.load("digits.pt"))
model.eval()   # ready to predict, no training needed

You did it — where to go next

In about half an hour you installed PyTorch and built a complete deep-learning project: data → model → training loop → evaluation → a saved model. That five-step shape is the same whether you're training a digit classifier or a large language model. Everything else is scale.

A few natural next steps to make it your own:

• Push the accuracy up. Add another Linear layer, widen 128 to 256, or train for more epochs. See how far a plain network gets.
• Go convolutional. Swap the linear layers for nn.Conv2d layers — the architecture built for images. Our From Pixels to Predictions course explains why they work.
• Understand the theory. The How a Network Actually Learns course opens up the .backward() step you leaned on, and How AI Learns places this project in the bigger map of machine learning.

You're no longer just reading about AI — you're training it.