Train the network

This is the heart of the project — and of all deep learning. Training is a loop that repeats one simple cycle over and over: guess, measure how wrong, nudge the weights to be a little less wrong. Do that a few thousand times and random noise turns into a digit reader.

We need two new tools before the loop: a way to measure wrongness (a loss function) and a way to apply the nudge (an optimizer).

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

• CrossEntropyLoss takes the model's 10 raw scores and the true label and returns a single number: how badly it missed. Confidently correct → near 0; confidently wrong → large. (It applies softmax internally, which is why the model outputs raw logits.)
• Adam is the optimizer. Given each weight's gradient, it decides how big a step to take. lr=1e-3 is the learning rate — the step size. A solid default.

The loop

Here it is — the four lines that do the actual learning, wrapped in a loop over the data. Add this to train.py:

epochs = 3
for epoch in range(epochs):
  model.train()
  for images, labels in train_loader:
      images, labels = images.to(device), labels.to(device)

      preds = model(images)            # 1. forward pass: guess
      loss = loss_fn(preds, labels)    # 2. how wrong were we?

      optimizer.zero_grad()            # 3a. clear old gradients
      loss.backward()                  # 3b. autograd: gradient per weight
      optimizer.step()                 # 3c. nudge every weight downhill

  print(f"epoch {epoch + 1}: loss {loss.item():.4f}")

Walk through the four lines that matter — this is the cycle from the tensor chapter, now at full scale:

1. model(images) — the forward pass. Push a batch through the network to get scores.
2. loss_fn(preds, labels) — compare those scores to the truth; get one wrongness number.
3. loss.backward() — the same .backward() you met in Chapter 2, now computing the gradient for every weight in the network at once. zero_grad() first, because PyTorch accumulates gradients by default and we want a fresh measurement each batch.
4. optimizer.step() — take one small step down each gradient. The weights get slightly better.

An epoch is one full pass over all 60,000 images. We do three.

Run it

python train.py

epoch 1: loss 0.2043
epoch 2: loss 0.1179
epoch 3: loss 0.0872

The exact numbers will vary, but the story won't: the loss falls. That downward drift is the network learning — each epoch, its guesses get closer to the truth. On a CPU this takes well under a minute. You've just trained a neural network. Last step: find out how good it actually is.