A PyTorch training loop that is not learning is one of the most common and frustrating situations in deep learning. The loss is flat, or NaN, or oscillating — and the model is right, the data is right, so what is wrong? This guide covers the systematic debugging process: how to isolate whether the problem is in the data, the model, the loss function, or the optimiser, and the specific checks that catch the most common causes.
Step 1: Overfit a Single Batch
The first test for any new training loop: can it memorise a single batch? If it cannot, the loop has a bug. Genuine learning difficulty only matters once you know the mechanics work correctly.
import torch, torch.nn as nn
# Pull one batch and train on it until loss → 0
model.train()
x, y = next(iter(train_loader))
x, y = x.to(device), y.to(device)
for step in range(500):
optimizer.zero_grad()
out = model(x)
loss = criterion(out, y)
loss.backward()
optimizer.step()
if step % 50 == 0:
print(f"step {step}: loss={loss.item():.6f}")
Expected: loss should reach near-zero (or at least drop sharply) within a few hundred steps on a small batch. If it does not, the issue is in the model or loss — not data volume or learning rate. Fix this before running on the full dataset.
Step 2: Check for NaN and Inf
NaN loss is almost always caused by one of: exploding gradients, a log of zero, division by zero in a custom loss, or bad input data. Add checks to isolate where the NaN first appears:
def check_nan(name, tensor):
if torch.isnan(tensor).any():
print(f"NaN in {name}")
if torch.isinf(tensor).any():
print(f"Inf in {name}")
# In your training step:
check_nan('input', x)
out = model(x)
check_nan('output', out)
loss = criterion(out, y)
check_nan('loss', loss)
loss.backward()
for name, param in model.named_parameters():
if param.grad is not None:
check_nan(f'grad/{name}', param.grad)
Enable PyTorch’s anomaly detection to get a full stack trace pointing to the exact operation that produced the NaN:
torch.autograd.set_detect_anomaly(True)
# run a few steps — anomaly detection is slow, disable after finding the bug
Step 3: Gradient Checks
If the model is not learning, gradients may be zero (vanishing) or very large (exploding). Check gradient norms after backward:
loss.backward()
total_norm = 0.0
zero_grad_params = []
for name, param in model.named_parameters():
if param.grad is None:
zero_grad_params.append(name)
continue
param_norm = param.grad.data.norm(2)
total_norm += param_norm.item() ** 2
if param_norm < 1e-8:
print(f"Near-zero grad: {name} norm={param_norm:.2e}")
total_norm = total_norm ** 0.5
print(f"Total grad norm: {total_norm:.4f}")
if zero_grad_params:
print(f"No gradient: {zero_grad_params}")
Total gradient norm consistently above 10 suggests exploding gradients — add gradient clipping: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) before optimizer.step(). Total norm consistently below 1e-4 suggests vanishing gradients — check activation functions (ReLU dying neurons), initialisation, or skip connections.
Step 4: Verify the Loss Function
Loss function mismatches are a frequent silent bug. The most common: using CrossEntropyLoss with softmax already applied in the model (double-softmax kills gradients), using NLLLoss without log-softmax, or using the wrong reduction.
# CrossEntropyLoss expects raw logits — do NOT apply softmax before it
criterion = nn.CrossEntropyLoss() # correct
out = model(x) # should be logits, not probabilities
loss = criterion(out, y)
# BCELoss expects probabilities in [0,1] — apply sigmoid first
criterion = nn.BCELoss()
out = torch.sigmoid(model(x)) # correct
# OR use BCEWithLogitsLoss (preferred — more numerically stable)
criterion = nn.BCEWithLogitsLoss()
out = model(x) # raw logits
# Check that targets match what the loss expects
print(f"y dtype: {y.dtype}, shape: {y.shape}")
# CrossEntropyLoss: y should be Long (int64), shape (batch,)
# MSELoss: y should be Float, same shape as output
Step 5: Check Data and Labels
Data bugs are common and easy to miss. Run these checks before training:
import matplotlib.pyplot as plt
x, y = next(iter(train_loader))
# Check shapes
print(f"x: {x.shape}, y: {y.shape}")
# Check value ranges
print(f"x range: [{x.min():.3f}, {x.max():.3f}], mean: {x.mean():.3f}")
print(f"y unique values: {y.unique()}, dtype: {y.dtype}")
# Check for class imbalance
from collections import Counter
all_labels = []
for _, labels in train_loader:
all_labels.extend(labels.tolist())
print(Counter(all_labels))
# Visualise a batch
fig, axes = plt.subplots(4, 8, figsize=(16, 8))
for i, ax in enumerate(axes.flat):
if i < x.shape[0]:
ax.imshow(x[i].squeeze(), cmap='gray')
ax.set_title(str(y[i].item()))
ax.axis('off')
plt.tight_layout(); plt.show()
Figure 1 — PyTorch training loop debugging checklist
Step 6: Learning Rate Diagnosis
Learning rate is often the issue when the model partially learns but plateaus or oscillates. Run a learning rate range test before committing to a fixed LR:
from torch.optim.lr_scheduler import OneCycleLR
import matplotlib.pyplot as plt
# LR range test: increase LR exponentially over a few hundred batches
losses, lrs = [], []
lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=1.05)
model.train()
for i, (x, y) in enumerate(train_loader):
if i > 200: break
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
lr_scheduler.step()
losses.append(loss.item())
lrs.append(optimizer.param_groups[0]['lr'])
plt.plot(lrs, losses); plt.xscale('log')
plt.xlabel('LR'); plt.ylabel('Loss')
plt.title('LR Range Test — pick LR at steepest decline')
plt.show()
Step 7: Check the Optimiser State
Two common optimiser bugs. Calling optimizer.zero_grad() in the wrong place (after backward instead of before forward) accumulates gradients across steps. Not calling optimizer.step() at all silently prevents weight updates — the model receives gradients but parameters never change. Verify parameter values are actually changing:
param_before = next(model.parameters()).data.clone()
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
param_after = next(model.parameters()).data
diff = (param_after - param_before).abs().max()
print(f"Max param change: {diff:.6f}")
# Should be > 0. If 0, parameters are not updating.
Step 8: train() vs eval() Mode
Calling model.eval() during training disables dropout and freezes BatchNorm statistics. If you forget to switch back to model.train() after validation, training silently runs in eval mode — a hard-to-spot bug that makes the model appear to learn slowly or not at all in models that rely on dropout regularisation:
for epoch in range(n_epochs):
model.train() # ← required before training loop
for x, y in train_loader:
...
model.eval() # ← required before validation
with torch.no_grad():
for x, y in val_loader:
...
# model.train() will be called at the top of the next epoch iteration
The single-batch overfit test catches most training loop bugs in under five minutes. Run it first on every new architecture before doing anything else. If the loop can memorise one batch, the mechanics are correct — then the problem is scale, regularisation, or hyperparameters, which are much easier to diagnose. If it cannot memorise one batch, work through the gradient, loss, and data checks above until it can, then move on to full training.
Figure 1 — PyTorch training loop debugging checklist
Profiling Slow Training
When training runs but is unexpectedly slow, use PyTorch's built-in profiler to find the bottleneck — usually DataLoader I/O, CPU-GPU transfer, or an inefficient operation:
from torch.profiler import profile, record_function, ProfilerActivity
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True
) as prof:
for i, (x, y) in enumerate(train_loader):
if i == 10: break
with record_function("data_transfer"):
x, y = x.to(device), y.to(device)
with record_function("forward"):
out = model(x)
with record_function("loss"):
loss = criterion(out, y)
with record_function("backward"):
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))Look for operations that take disproportionate time. If data_transfer dominates, use pin_memory=True and non_blocking=True in your DataLoader and .to(device) calls. If a single layer dominates, consider a more efficient implementation or check for unnecessary CPU-GPU synchronisation points.
Reproducibility Checklist
If your training results vary across runs, lock down the random seeds:
import torch, numpy as np, random, os
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
os.environ['PYTHONHASHSEED'] = str(seed)
set_seed(42)Note that cudnn.deterministic = True disables some optimisations and may slow training by 5–15%. For debugging, the reproducibility is worth it. For production training, set it only if bit-exact reproducibility is required — otherwise leave benchmark=True for speed.
Verifying Validation Logic
Validation bugs are subtle because they do not prevent training — they just produce misleading metrics. Common issues: computing accuracy on logits instead of predictions (out.argmax(1)), accumulating metrics incorrectly across batches (averaging averages instead of summing then dividing), or inadvertently including gradient computation in validation. A quick sanity check — does your validation accuracy on the training set match training accuracy? If validation is run correctly on training data, the numbers should be close. If they diverge significantly, the validation loop has a bug.
def evaluate(model, loader, device):
model.eval()
correct, total = 0, 0
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
preds = model(x).argmax(dim=1) # argmax on logits
correct += (preds == y).sum().item()
total += y.size(0)
return correct / total # accuracy over ALL examples
# Sanity check: eval on train set should match train accuracy
train_acc = evaluate(model, train_loader, device)
val_acc = evaluate(model, val_loader, device)
print(f"Train acc: {train_acc:.4f} Val acc: {val_acc:.4f}")Memory Leak Detection
If GPU memory grows monotonically during training, you have a memory leak — usually caused by storing tensors that retain their computation graph. The most common source: accumulating losses for logging without detaching them.
# Bad — retains computation graph, causes memory leak
total_loss += loss
# Good — detach before accumulating
total_loss += loss.item() # .item() returns a Python float
# or
total_loss += loss.detach().cpu() # if you need a tensor
# Monitor GPU memory during training
for i, (x, y) in enumerate(train_loader):
...
if i % 50 == 0:
mem = torch.cuda.memory_allocated() / 1e9
print(f"Step {i}: GPU mem {mem:.2f} GB")
# Should be roughly constant — if growing, you have a leakCommon Bugs Quick Reference
The bugs that account for the majority of non-learning training loops. Forgetting zero_grad() — gradients accumulate across steps, producing wildly large updates that diverge the loss. Always call optimizer.zero_grad() at the start of each iteration. Double softmax — applying torch.softmax() inside the model and then passing to CrossEntropyLoss (which applies log-softmax internally) squashes gradients toward zero; remove the explicit softmax. Target dtype mismatch — CrossEntropyLoss requires Long (int64) targets; passing Float silently produces wrong results or an error depending on PyTorch version. Model in eval() during training — dropout is disabled and BatchNorm runs on stored statistics rather than batch statistics; the model may appear to train but generalisation is wrong. Not moving data to device — a CUDA tensor model receives CPU data and raises an error; always call x, y = x.to(device), y.to(device) at the start of each step. Loss reduction mismatch — reduction='sum' vs reduction='mean' changes the gradient scale with batch size; use mean unless you have a specific reason for sum.
Using torch.autograd.gradcheck
For custom loss functions or autograd operations, use gradcheck to verify the analytical gradient against a numerical finite-difference approximation:
from torch.autograd import gradcheck
def my_loss(input, target):
return ((input - target) ** 2).sum()
input = torch.randn(4, 10, dtype=torch.float64, requires_grad=True)
target = torch.randn(4, 10, dtype=torch.float64)
# gradcheck requires float64 for numerical precision
result = gradcheck(my_loss, (input, target), eps=1e-6, atol=1e-4)
print(f"gradcheck passed: {result}")A failed gradcheck means your custom operation's backward pass is incorrect — the gradient you compute does not match what the loss actually changes. This is the definitive test for custom autograd functions.
Systematising Your Debugging Workflow
The most efficient debugging workflow: keep a short checklist and run it in order. Do not skip steps because "it's probably not that" — the overfit test takes two minutes and has caught bugs that experienced practitioners spent hours hunting by other means. Once you have run this checklist a few times on different projects, the sequence becomes automatic and most bugs are caught in under ten minutes. Add project-specific checks at the end — for example, checking that attention masks are applied correctly in a transformer, or that the KL divergence term in a VAE is not being divided by the wrong batch dimension. The generic checks above cover the universal failure modes; project-specific checks cover the rest.
The single-batch overfit test catches most training loop bugs in under five minutes. Run it first on every new architecture before doing anything else. If the loop can memorise one batch, the mechanics are correct — then the problem is scale, regularisation, or hyperparameters, which are much easier to diagnose. Work through the steps above in order; most bugs surface by Step 3 or 4, and rarely do you need to go beyond Step 6 to find the issue. The two most time-saving habits: always run the single-batch overfit test before any multi-epoch run, and always check gradient norms for the first few steps of a new architecture — these two checks alone catch the vast majority of training loop bugs before they waste hours of GPU time.
Debugging a training loop is a skill that improves with practice — each bug you track down systematically makes the next one faster to find, because the patterns repeat across projects and architectures.