Exercise — Debug a Training Loop

Build the mental model

When a PyTorch training loop runs without errors but the loss refuses to drop, resist the urge to guess — work through a fixed checklist in order, since each check rules out an entire class of failure. First: is optimizer.zero_grad() called every iteration, before backward()? Gradients accumulate by default, so a missing zero_grad() silently sums gradients across batches, producing update directions that make no sense after the first few steps. Second: is loss.backward() actually invoked before optimizer.step()? Skipping it leaves .grad at its previous value (often None), and step() either does nothing or reuses stale gradients. Third: is the model in the correct mode — model.train() before the loop, not left in .eval()? Eval mode freezes dropout and batch-norm statistics, quietly capping learning capacity without raising an error. Fourth: is the learning rate wildly too high (loss oscillates or turns NaN) or too low (loss inches down over hundreds of epochs)? Fifth: does the loss function match the task — MSELoss on class logits instead of CrossEntropyLoss trains toward the wrong objective entirely, which looks like 'not learning' but is really 'learning the wrong thing'. Work top to bottom before touching hyperparameters.

Connect it to a real scenario

Suppose Thuta Learning's engineering team ships a sentiment classifier meant to flag frustrated learners in course-feedback text, but after a week of training the model still predicts the same neutral label for everything. Before assuming the architecture is wrong, you're asked to audit the training script using exactly this checklist: confirm zero_grad() runs every batch, confirm backward() precedes step(), confirm the model is in .train() mode during training and only switched to .eval() for validation, sanity-check the learning rate against the loss curve, and verify CrossEntropyLoss — not MSELoss — is paired with raw class-index labels. This audit habit catches the majority of 'trains but doesn't learn' bugs before anyone touches the model design.

• A flat or non-decreasing loss with no crash almost always traces to one of five checks: zero_grad(), backward(), train()/eval() mode, learning rate scale, or loss-function choice.
• Missing optimizer.zero_grad() doesn't error — it silently accumulates gradients across batches, which is far more dangerous than a crash because nothing signals the bug.
• Leaving a model in .eval() mode during training is a common copy-paste mistake that disables dropout and batch-norm updates without any warning.

Try the working example

import torch
import torch.nn as nn
import torch.optim as optim

class FeedbackClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 10),
        )

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

model = FeedbackClassifier()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

model.eval()  # set up before training starts

for epoch in range(10):
    for inputs, labels in train_loader:
        inputs = inputs.view(inputs.size(0), -1)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch}: loss = {loss.item():.4f}")

Running the loop as written prints something like `Epoch 0: loss = 0.9987`, `Epoch 1: loss = 0.9991`, `Epoch 9: loss = 0.9979` — the loss barely moves at all, which is the symptom, not the diagnosis; the buggy lines above are the cause.

PyTorch Docs — Optimization — Deep Learning