Debugging training: shapes, NaNs and silent failures

The usual suspects

Shape errors are most common. mat1 and mat2 shapes cannot be multiplied (32x3136 and 2048x10) means your input reaches a Linear layer with 3136 features but the layer expects 2048; fix in_features or the layers before it. Print x.shape inside forward until the numbers make sense.

Dtype errors: expected scalar type Float but found Double (NumPy float64 data; add .float()), or a cross-entropy target that is float instead of int64.

Loss becomes NaN: learning rate too high, log(0) in a hand-written loss, or bad values in the data. torch.autograd.set_detect_anomaly(True) reports which backward operation first produced NaN, but it is slow.

Loss does not go down is the hardest, because nothing crashes. Check for a missing zero_grad() or step(), softmax before CrossEntropyLoss, a frozen model, labels misaligned with inputs, a learning rate far too small, or the model left in eval mode. The fastest test: try to overfit a single small batch. A correct model and loss can memorise it to near-zero loss; if yours cannot, the bug is in the code, not in the amount of data.

Overfit one batch

Take one batch and train on it repeatedly. The loss should fall towards zero within a few hundred steps. train_loader, model, loss_fn, optimizer and device come from your training script.

import torch

xb, yb = next(iter(train_loader))
xb, yb = xb.to(device), yb.to(device)
model.train()
for step in range(300):
    optimizer.zero_grad()
    loss = loss_fn(model(xb), yb)
    loss.backward()
    optimizer.step()
    if step % 50 == 0:
        print(step, round(loss.item(), 4))
# Expect the loss to approach 0. If it does not, debug the model and loss first.

Test the recipe on one plate

Before cooking biryani for five hundred wedding guests, a caterer makes one plate and tastes it. If one plate is wrong, cooking more will not fix it. Overfitting a single batch is that one plate: if the model cannot get even a few examples right, more data and longer training only waste GPU hours.

A quick checklist before long runs

Visualise a few input samples with their labels after all transforms. Check the first loss against ln(number of classes). Overfit one batch. Confirm model.parameters() count looks right. Watch both training and validation loss: training falling while validation rises means overfitting, not a bug.

त्वरित जाँच: Your training loss stays flat from the first epoch. Which test best separates a code bug from a data-size problem?

• Try to overfit a single small batch to near-zero loss
• Double the number of epochs
• Switch to a larger pretrained model
• Add more data augmentation