Top 5 PyTorch Tips to Debug Your Models Without Losing Your Mind

Manicandane PARI

— 5 min read
Share

When you start using PyTorch, it quickly feels like you’re fighting the framework:

  • Shapes don’t match
  • The loss doesn’t move
  • The GPU complains
  • Error messages look cryptic

The good news: 80% of the problems come from the same few causes.
If you get a few simple habits right, you can debug most models in minutes, not hours.

In this article, we’ll go through 5 practical PyTorch tips (with copy-paste code) to help you find what’s wrong faster.

1. Check tensor shapes at every key step

Most PyTorch issues are just shape mismatches.
If you’ve ever seen:

RuntimeError: mat1 and mat2 shapes cannot be multiplied

…you know exactly what I mean.

Habit: print .shape at critical points

Whenever you do view, reshape, permute, cat, etc., add one debug line:

  • at the beginning of forward
  • after flattening
  • after concatenations
  • before the final layer

Example with a simple network:

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        print("Input shape :", x.shape)      # (batch_size, 1, 28, 28) ?
        x = x.view(x.size(0), -1)           # flatten
        print("After view  :", x.shape)     # (batch_size, 784)
        x = self.fc1(x)
        print("After fc1   :", x.shape)     # (batch_size, 128)
        x = self.fc2(x)
        print("After fc2   :", x.shape)     # (batch_size, 10)
        return x

You can enforce expectations with assertions:

def forward(self, x):
    # We expect (batch_size, 1, 28, 28)
    assert x.ndim == 4 and x.shape[1:] == (1, 28, 28), f"Unexpected shape: {x.shape}"
    x = x.view(x.size(0), -1)
    assert x.shape[1] == 28 * 28, f"Incorrect flatten shape: {x.shape}"
    ...

💡 Mental rule: if you’re lost, print shapes from raw data → final logits, like tracing a pipeline.

2. Try to overfit a tiny batch on purpose

This single trick can save you a lot of time.

If your model can’t overfit 32 examples, it won’t work on 32,000.

The goal: check that model + loss + labels are coherent.

Procedure

  1. Take one batch from your DataLoader (e.g., 32 samples).
  2. Train your model only on that batch for many epochs.
  3. Check if the loss goes close to zero.

If the loss stays high → something is fundamentally wrong.

“Overfit a batch” example

model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# Get ONE batch
data_iter = iter(train_loader)
inputs, targets = next(data_iter)

for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        print(f"Epoch {epoch}, loss = {loss.item():.4f}")

How to interpret this

  • If the loss drops a lot and becomes very small → your model can learn.
    The problem might be elsewhere (LR, regularization, noisy data…).

  • If the loss barely moves:

    • there might be a bug in the forward
    • labels may be wrong (wrong type/shape, misaligned)
    • you may be using the wrong loss (e.g. MSELoss with class indices)
    • your learning rate is completely off

💡 Rule of thumb: as long as you cannot overfit a small batch, don’t bother launching a 10-hour training run.

3. Look at the loss… but also at some predictions

We often stare only at the loss curve, but it’s very abstract.
Sometimes the model is actually learning, but you don’t realize it from the loss alone.

Simple habit: print a few predictions vs labels every epoch (or every N epochs).

Minimal example

for epoch in range(num_epochs):
    model.train()
    for inputs, targets in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

    # Quick inspection at the end of the epoch
    model.eval()
    with torch.no_grad():
        sample_inputs, sample_targets = next(iter(val_loader))
        sample_outputs = model(sample_inputs)
        preds = sample_outputs.argmax(dim=1)

        print(f"\nEpoch {epoch} - train loss: {loss.item():.4f}")
        print("Targets :", sample_targets[:10].tolist())
        print("Preds   :", preds[:10].tolist())

What to watch for

  • Do predictions change over epochs?
  • Does the model always output the same class (e.g., [3, 3, 3, 3, 3, 3, ...])?
  • Do predictions gradually become more consistent with labels?

If predictions are random or frozen:

  • your learning rate might be too low → everything moves too slowly
  • or too high → training diverges
  • or you have a target bug (wrong label mapping, shifted indices, etc.)

💡 Even without fancy plots, just printing 10 predictions and labels each epoch already gives you a huge amount of signal.

4. Hunt down device and dtype issues

Two ultra-common categories of errors:

  1. Mixing CPU and GPU:

    RuntimeError: Expected all tensors to be on the same device, but found at least two devices

  2. Wrong type for the loss:

    RuntimeError: expected scalar type Long but found Float

Base pattern you should always use

Always start by defining a device:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = SimpleNet().to(device)
criterion = nn.CrossEntropyLoss()

for inputs, targets in train_loader:
    inputs = inputs.to(device)
    targets = targets.to(device)  # important: same device as inputs & model

    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

If you’re unsure, log the devices:

print(inputs.device, targets.device, next(model.parameters()).device)

Watch dtype when using losses

Classic example:

  • For nn.CrossEntropyLoss:

    • outputs (logits) → torch.float32
    • targets (class indices) → torch.int64 (LongTensor)

You can check:

print(outputs.dtype)  # torch.float32
print(targets.dtype)  # torch.int64 expected for CrossEntropyLoss

If your targets are floats (for example after some accidental normalization or cast), things will break or behave weirdly.

💡 Quick debug order when you see a weird error:

  1. Check devices (CPU vs GPU).
  2. Check dtypes (float vs long, etc.).
  3. Only then dig deeper.

5. Isolate the problem with a simple forward and dummy data

Before any big training run, you should check that your model can do a forward pass without crashing.
You can do this with dummy data (torch.randn) that has the same shape as real data.

Test 1: forward with random inputs

batch_size = 16
dummy_inputs = torch.randn(batch_size, 1, 28, 28)  # same shape as your real images

model = SimpleNet()
dummy_outputs = model(dummy_inputs)
print("Dummy output shape :", dummy_outputs.shape)  # expected: (16, 10)

If it crashes here, you know:

  • the bug is in the model architecture (dimensions, reshapes, etc.)
  • not in the DataLoader, optimizer, or loss.

Test 2: forward in eval mode, no gradients

model.eval()
with torch.no_grad():
    dummy_outputs = model(dummy_inputs)
    print(dummy_outputs[:5])

Then, do the same with a real batch from your DataLoader:

model.eval()
with torch.no_grad():
    real_inputs, real_targets = next(iter(train_loader))
    real_outputs = model(real_inputs)
    print("Real output shape :", real_outputs.shape)

If dummy inputs work but real ones don’t:

  • the problem is probably in data preparation (resizing, normalization, channels order, etc.).

💡 Ideal routine before any long training:

  1. Forward with torch.randn → OK?
  2. Forward with one real batch → OK?
  3. Overfit one small batch → OK?
  4. Then launch full training.

Wrap-up: 5 habits for every PyTorch project

To summarize, most PyTorch pain can be reduced with a few simple habits:

  1. Check tensor shapes at key points (input, flatten, concat, output).
  2. Try to overfit a tiny batch to validate model + loss + labels.
  3. Inspect a few predictions each epoch, not just the loss value.
  4. Verify device (CPU/GPU) and dtype before you overthink the bug.
  5. Test the forward pass with dummy data, then with a real batch, before long runs.

If you turn these five points into a mental checklist, you’ll spend much less time on “stupid bugs” and much more time on building actually useful models.

Read more