# Self-check: spot the bugs — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-spot-the-bug

> Read a realistic but broken training script, find each bug, and explain the symptom each one would cause.

## 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.

```python
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()`.

**Quiz:** 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] `x.to("cuda")` without reassigning `x`
- [ ] The softmax before `CrossEntropyLoss`

*Answer:* `x.to("cuda")` without reassigning `x`. `x` stays on the CPU while the model is on the GPU, so the first forward pass fails with a device mismatch. The other bugs run silently and only hurt results.