Define the network

Now the fun part: the neural network itself. In PyTorch a model is a Python class that subclasses nn.Module. You describe its layers in __init__, and how data flows through them in forward. That's the whole contract.

Our network is deliberately simple — and it still hits ~97% accuracy on digits:

class DigitClassifier(nn.Module):
  def __init__(self):
      super().__init__()
      self.net = nn.Sequential(
          nn.Flatten(),            # 1×28×28 image → 784-long vector
          nn.Linear(28 * 28, 128), # fully-connected: 784 → 128
          nn.ReLU(),               # non-linearity
          nn.Linear(128, 10),      # 128 → 10 scores, one per digit
      )

  def forward(self, x):
      return self.net(x)

model = DigitClassifier().to(device)
print(model)

What each layer does

Read the Sequential block top to bottom — it's the exact path a batch of images takes:

• Flatten turns each 1×28×28 image into a flat row of 784 numbers. The network doesn't see a grid; it sees a list of pixel brightnesses. (This is the reshape move from Chapter 2, done for you.)
• Linear(784, 128) is a fully-connected layer: every one of the 784 inputs connects to each of 128 outputs, through weights the network will learn. This is where almost all the model's knowledge lives.
• ReLU ("rectified linear unit") replaces negatives with zero. Without a non-linear step like this between the linear layers, stacking them would collapse into a single linear layer — and no straight line can separate ten kinds of digit.
• Linear(128, 10) produces 10 scores, one per possible digit. The highest score is the network's guess.

It already runs — just badly

The model works right now; it's simply untrained, so its weights are random and its guesses are noise. Prove it by pushing one batch through:

images, labels = next(iter(train_loader))
scores = model(images.to(device))
print(scores.shape)          # torch.Size([64, 10]) — 10 scores per image
print(scores[0])             # ten raw numbers; argmax is the guess

Those ten numbers per image are called logits — raw, unnormalised scores. We don't need to turn them into probabilities ourselves; the loss function in the next chapter does that internally. Speaking of which: right now the network is guessing at random. Let's teach it.