PyTorch Neural Network Guide: 5 Mistakes That Break Training

Why Your First PyTorch Model Probably Won’t Train

Your loss curve flatlines at 0.693. The gradients are all NaN after epoch 2. Or worse — training completes without errors, but your model predicts the same class for every input.

I’ve seen these failure modes dozens of times, and they all trace back to the same handful of setup mistakes. PyTorch gives you enough rope to hang yourself: it won’t stop you from initializing weights incorrectly, forgetting to zero gradients, or using the wrong loss function for your task. The training loop runs, the progress bar fills up, and you only realize something’s wrong when you check the outputs.

This isn’t a gentle introduction to neural networks. It’s a focused look at the five PyTorch-specific mistakes that silently break training, why they happen, and how to fix them before you waste GPU hours. I’m assuming you know what backpropagation is and have written at least one forward() method. If you’re still Googling “what is a tensor,” start with the official PyTorch tutorials and come back.

The Gradient Accumulation Trap

Here’s the mistake: you write a training loop, it runs, loss goes down, everything looks fine. Then you try to add gradient accumulation for larger effective batch sizes, and suddenly your model diverges.

# This looks innocent but breaks training
model = SimpleNet()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(10):
    for i, (inputs, targets) in enumerate(dataloader):
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()

        # Only step every 4 batches for gradient accumulation
        if (i + 1) % 4 == 0:
            optimizer.step()
            optimizer.zero_grad()

The problem: you’re accumulating gradients for batches 0, 1, 2, 3 — then stepping. But what about the very last batch in the epoch if len(dataloader) % 4 != 0? Those gradients never get applied. You’ve effectively reduced your training set size.

Worse, if you forget zero_grad() entirely (I’ve done this), gradients keep piling up across epochs. Your effective learning rate becomes batch_size × num_batches × num_epochs × lr, which is astronomical. The model either explodes immediately or oscillates wildly. This debugging light has saved me during late-night gradient debugging sessions — stick it behind your monitor and set it to red when loss > 1.0, green when converging. Silly, but it works.

The fix:

accumulation_steps = 4
model.zero_grad()  # Zero once at the start

for epoch in range(10):
    for i, (inputs, targets) in enumerate(dataloader):
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss = loss / accumulation_steps  # Scale loss
        loss.backward()

        if (i + 1) % accumulation_steps == 0:
            optimizer.step()
            model.zero_grad()

    # Critical: step at epoch end for remaining gradients
    if (i + 1) % accumulation_steps != 0:
        optimizer.step()
        model.zero_grad()

Notice loss / accumulation_steps — this keeps gradient magnitudes consistent with non-accumulated training. Without it, your gradients are 4× larger than expected, which effectively multiplies your learning rate by 4.

Weight Initialization: When PyTorch Defaults Fail

PyTorch initializes nn.Linear layers with uniform Kaiming initialization by default, which works fine for ReLU activations. But swap to Tanh or Sigmoid, and you’ll hit saturation before training even starts.

import torch
import torch.nn as nn

class DeepNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(784, 256),
            nn.Tanh(),
            nn.Linear(256, 128),
            nn.Tanh(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        return self.layers(x)

model = DeepNet()
x = torch.randn(32, 784)  # Batch of 32 MNIST images

with torch.no_grad():
    activations = []
    for layer in model.layers:
        x = layer(x)
        if isinstance(layer, nn.Tanh):
            activations.append(x.std().item())

print(f"Activation stds: {activations}")  # [0.68, 0.23] — shrinking fast

On my M1 MacBook with PyTorch 2.2, those Tanh activations have std around 0.68 and 0.23. By the third Tanh layer (not shown), you’d be below 0.1. Most neurons are stuck near zero where $\tanh'(x) \approx 0$, so gradients vanish. The network barely learns.

Xavier initialization fixes this for Tanh:

class DeepNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(784, 256),
            nn.Tanh(),
            nn.Linear(256, 128),
            nn.Tanh(),
            nn.Linear(128, 10)
        )

        # Xavier init for Tanh
        for layer in self.layers:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_normal_(layer.weight)
                nn.init.zeros_(layer.bias)

Xavier (also called Glorot) initialization sets weights $\sim \mathcal{N}(0, \sqrt{2/(n_{\text{in}} + n_{\text{out}})})$, which maintains activation variance across layers for symmetric activations like Tanh. For ReLU, use Kaiming: $\sim \mathcal{N}(0, \sqrt{2/n_{\text{in}}})$.

Here’s the counterintuitive part: modern architectures like ResNets and Transformers often use nn.init.xavier_uniform_ even with ReLU, because skip connections and layer norm change the dynamics. The official PyTorch ResNet implementation uses Kaiming for conv layers but leaves linear layers at default. Vision Transformers (Dosovitskiy et al., 2021) use truncated normal for patch embeddings. It’s not one-size-fits-all.

When in doubt, check activation statistics after the forward pass on a random batch. If stds are collapsing toward zero or exploding past 2.0, your initialization is wrong.

Loss Function Mismatch: BCELoss vs BCEWithLogitsLoss

This one is subtle and crashes silently with NaN losses.

# Binary classification — which loss is correct?
class BinaryClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 1)
        # self.sigmoid = nn.Sigmoid()  # Do you need this?

    def forward(self, x):
        return self.fc(x)  # Logits or probabilities?

model = BinaryClassifier()
criterion = nn.BCELoss()  # Expects probabilities in [0, 1]
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

x = torch.randn(8, 10)
y = torch.randint(0, 2, (8, 1)).float()

preds = model(x)  # Output: raw logits, e.g. tensor([[-2.3], [1.5], ...])
loss = criterion(preds, y)  # BOOM: RuntimeError or NaN

The model outputs raw logits (unbounded), but nn.BCELoss expects probabilities in $[0, 1]$. If you pass negative logits to BCELoss, it either errors out or computes $\log(\text{negative number})$ = NaN.

The fix: use nn.BCEWithLogitsLoss, which combines sigmoid + BCE in one numerically stable operation.

criterion = nn.BCEWithLogitsLoss()
loss = criterion(preds, y)  # Works — applies sigmoid internally

Internally, BCEWithLogitsLoss computes:

$$L = -\frac{1}{N} \sum_{i=1}^{N} \left[ y_i \log(\sigma(z_i)) + (1 – y_i) \log(1 – \sigma(z_i)) \right]$$

where $\sigma(z) = 1/(1 + e^{-z})$. But it uses the log-sum-exp trick to avoid numerical underflow when $z_i$ is very negative. If you manually apply sigmoid then BCELoss, you lose that numerical safety.

Same principle for multi-class: use nn.CrossEntropyLoss (logits) not nn.NLLLoss (log-probabilities) unless you explicitly apply log_softmax. I wasted a full afternoon debugging a text classifier because I used NLLLoss with raw logits and got garbage predictions.

Learning Rate Schedulers That Destroy Convergence

You’ve probably seen torch.optim.lr_scheduler.ReduceLROnPlateau in tutorials — reduce LR when validation loss plateaus. Sounds great. But if you wire it up wrong, it can permanently cripple training.

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.1, patience=5
)

for epoch in range(50):
    train_loss = train_one_epoch(model, train_loader, optimizer)
    val_loss = validate(model, val_loader)

    # WRONG: stepping on train_loss
    scheduler.step(train_loss)

If you step the scheduler on training loss instead of validation loss, it triggers LR reduction as soon as training loss stops improving — which happens naturally when the model fits the training set. Your LR drops from 0.001 to 0.0001 at epoch 10, then 0.00001 at epoch 15, and by epoch 20 you’re at 1e-6. Training effectively stops.

Always step on validation loss for ReduceLROnPlateau:

scheduler.step(val_loss)  # Correct

Another gotcha: CosineAnnealingLR with warm restarts. If you set T_max incorrectly, the LR might restart mid-training when you don’t want it to.

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=50, eta_min=1e-6
)

# If you train for 100 epochs, LR restarts at epoch 50
# Might disrupt late-stage convergence

The cosine schedule is:

$$\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} – \eta_{\min}) \left(1 + \cos\left(\frac{T_{\text{cur}}}{T_{\max}} \pi\right)\right)$$

where $T_{\text{cur}}$ is the current epoch and $T_{\max}$ is the cycle length. If your training is 100 epochs, set T_max=100 to avoid mid-training restarts.

I’ve also seen people call scheduler.step() inside the batch loop instead of once per epoch, which makes the LR decay 1000× faster than intended if you have 1000 batches. Read the docs carefully — some schedulers expect per-epoch stepping, others expect per-batch.

Device Mismatch: The Silent Killer

This is the most frustrating because PyTorch doesn’t always error immediately.

model = SimpleNet().to('cuda')
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for inputs, targets in dataloader:
    # inputs and targets are on CPU by default
    outputs = model(inputs)  # RuntimeError: Expected all tensors to be on the same device

Fine, you add .to('cuda'):

for inputs, targets in dataloader:
    inputs = inputs.to('cuda')
    targets = targets.to('cuda')
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

But now training is slow because you’re transferring data to GPU every batch. If you’re using a large batch size on a dataset that fits in VRAM, you’re burning 20-30% of training time on transfers. Better:

# Pin memory and non_blocking transfers
train_loader = DataLoader(
    dataset, batch_size=64, pin_memory=True, num_workers=4
)

for inputs, targets in train_loader:
    inputs = inputs.to('cuda', non_blocking=True)
    targets = targets.to('cuda', non_blocking=True)
    # ...

pin_memory=True allocates data in pinned (page-locked) RAM, which enables faster GPU transfers. non_blocking=True makes the transfer asynchronous — the CPU can start loading the next batch while the current one is copying.

But here’s the silent killer: custom layers that hardcode device.

class BrokenLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = torch.randn(10, 10)  # Always on CPU

    def forward(self, x):
        return x @ self.weight  # x is on GPU, weight is CPU → crash

The fix: register tensors as parameters or buffers.

class FixedLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('weight', torch.randn(10, 10))

    def forward(self, x):
        return x @ self.weight  # .to('cuda') moves buffer too

register_buffer tells PyTorch “this tensor is part of the module state, move it when I call .to(device).” Use it for anything that’s not a trainable parameter but needs to follow the model’s device.

When Debugging Fails: Gradient Checking

Sometimes training fails and you can’t figure out why. Gradients look fine, loss is computed correctly, but the model doesn’t learn. I fall back to numerical gradient checking — compare PyTorch’s autograd against finite differences.

def numerical_gradient(model, x, y, criterion, epsilon=1e-5):
    """Compute gradients via finite differences."""
    numerical_grads = {}

    for name, param in model.named_parameters():
        grad = torch.zeros_like(param)
        flat_param = param.data.flatten()

        for i in range(len(flat_param)):
            old_val = flat_param[i].item()

            # f(x + epsilon)
            flat_param[i] = old_val + epsilon
            loss_plus = criterion(model(x), y).item()

            # f(x - epsilon)
            flat_param[i] = old_val - epsilon
            loss_minus = criterion(model(x), y).item()

            # Gradient: (f(x+eps) - f(x-eps)) / 2*eps
            grad.flatten()[i] = (loss_plus - loss_minus) / (2 * epsilon)

            # Restore original value
            flat_param[i] = old_val

        numerical_grads[name] = grad

    return numerical_grads

# Compare against autograd
model = SimpleNet()
x = torch.randn(4, 10)
y = torch.randint(0, 2, (4,))
criterion = nn.CrossEntropyLoss()

loss = criterion(model(x), y)
loss.backward()

num_grads = numerical_gradient(model, x, y, criterion)

for name, param in model.named_parameters():
    autograd_grad = param.grad
    numerical_grad = num_grads[name]
    diff = (autograd_grad - numerical_grad).abs().max().item()
    print(f"{name}: max difference = {diff:.2e}")
    # Should be < 1e-5 for float32

If the difference is > 1e-3, something’s broken in your backward pass (usually a custom layer). This saved me when I implemented a custom attention mechanism and forgot to handle the case where attention weights sum to zero.

But numerical gradient checking is SLOW — it requires O(num_parameters) forward passes. Only use it on tiny models/batches for debugging.

Batch Normalization Inference Mistake

Batch norm behaves differently in training vs evaluation mode, and forgetting to switch modes breaks everything.

model = nn.Sequential(
    nn.Linear(10, 20),
    nn.BatchNorm1d(20),
    nn.ReLU(),
    nn.Linear(20, 2)
)

# Training
model.train()
for x, y in train_loader:
    loss = criterion(model(x), y)
    loss.backward()
    optimizer.step()

# Validation — FORGOT model.eval()
val_loss = 0
for x, y in val_loader:
    val_loss += criterion(model(x), y).item()  # Still in train mode!

In training mode, batch norm normalizes using batch statistics:

$$\hat{x}_i = \frac{x_i – \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}$$

where $\mu_B$ and $\sigma_B^2$ are the mean and variance of the current batch. In eval mode, it uses running statistics accumulated during training. If you forget model.eval(), validation loss is computed using batch stats, which are noisy for small eval batches. Your reported validation loss is artificially high, and early stopping might trigger too soon.

Always:

model.eval()
with torch.no_grad():
    for x, y in val_loader:
        val_loss += criterion(model(x), y).item()

Same goes for dropout — it’s active in train mode, disabled in eval mode. If you test a model with dropout still on, predictions are stochastic and accuracy is lower than it should be.

FAQ

Q: Should I use .detach() or .item() when logging loss?

Use .item() for scalar losses to avoid memory leaks. If you store the loss tensor itself in a list (e.g., losses.append(loss)), PyTorch keeps the entire computation graph in memory for the whole epoch. Use losses.append(loss.item()) to store just the float value.

Q: Why does my model train fine on GPU but crash on CPU with “illegal instruction”?

Probably a PyTorch build mismatch. If your CPU doesn’t support AVX2 but your PyTorch build assumes it does, you’ll get SIGILL. Reinstall PyTorch with a CPU-only build from pytorch.org that matches your hardware. I hit this on an old AWS c4 instance once.

Q: Is mixed precision training (AMP) always faster?

No. On older GPUs (pre-Volta, no Tensor Cores), AMP can actually be slower due to conversion overhead. On V100/A100/RTX 3090, it’s usually 1.5-2× faster with no accuracy loss. On CPU or ancient GPUs, stick to float32. And watch for training instabilities — some models need gradient scaling adjustments with AMP enabled.

What Actually Matters

If you take one thing from this: check your training loop for silent failures. PyTorch won’t stop you from writing code that runs but doesn’t train. Print gradient norms, log learning rates, visualize activation distributions, compare numerical vs autograd gradients. Most bugs hide in the places you assume are fine.

Use BCEWithLogitsLoss unless you have a very good reason not to. Zero gradients explicitly before backward(), not after step(). Set T_max correctly for cosine schedules. Call model.eval() before validation.

The patterns here generalize beyond PyTorch — JAX has similar gradient accumulation gotchas, TensorFlow’s tf.GradientTape requires manual zero-grad handling, and every framework has device placement surprises. But PyTorch’s flexibility makes these mistakes easier to make and harder to notice.

I’m still not entirely sure why PyTorch defaults to Kaiming init for all activations when Xavier is clearly better for Tanh/Sigmoid. My best guess is they optimized for the ReLU-dominant vision community. If anyone from the PyTorch team reads this, I’d love to know the design rationale.