Self-check: spot the bugs

How to read code for bugs

Debugging a training script is a skill you can practise by reading. Go through the code with four questions. Shapes and dtypes: what is the shape and dtype of every tensor at every line, and do the loss inputs match what the loss expects? Devices: is every tensor in each operation on the same device? Gradient bookkeeping: are gradients cleared every step, is backward() before step(), and is nothing needed for training computed under no_grad()? Modes: is the model in train() while training and eval() while evaluating?

Many bugs produce no error message. A softmax before cross-entropy, a missing zero_grad() or validation without eval() all run happily and just give worse numbers. That is why reading the code carefully, and testing with an overfit-one-batch run, beats waiting for a crash. Try the script below before reading the answers in the tip.

A script with five bugs

Assume train_ds and val_loader exist and the model is on a GPU. Find five problems before reading on.

import torch
from torch import nn
from torch.utils.data import DataLoader

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 3), nn.Softmax(dim=1))
model.to("cuda")
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loader = DataLoader(train_ds, batch_size=32, shuffle=False)

for epoch in range(5):
    for x, y in loader:
        x.to("cuda")
        y = y.to("cuda")
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()
    model.eval()
    # ... validation using val_loader ...

Answers

1. nn.Softmax before CrossEntropyLoss: remove it and feed logits. 2. shuffle=False for training: use shuffle=True. 3. x.to("cuda") is not reassigned: write x = x.to("cuda"), or the forward pass hits a device mismatch. 4. No optimizer.zero_grad(): gradients accumulate across steps. 5. model.eval() is called but model.train() never is, so after epoch 0 the model trains in eval mode (it matters once you add dropout or batch norm). Validation should also run under torch.inference_mode().

Quick check: In the script above, which bug would raise an exception on the very first batch?

• The missing `optimizer.zero_grad()`
• `shuffle=False` on the training loader
• `x.to("cuda")` without reassigning `x`
• The softmax before `CrossEntropyLoss`