PyTorch Training Loop — Missing zero_grad Causes Nonsense

The training loop is the core execution pattern in PyTorch — a cycle that repeats for every batch of data: clear gradients, compute predictions, compute loss, backpropagate, update weights. PyTorch keeps this explicit on purpose. You are never far from the mechanics of optimization.

That explicitness is the trade-off. You get full visibility into gradient flow, parameter updates, device placement, mixed precision, gradient clipping, and scheduling. The cost is that there is no place to hide sloppy thinking. The loop is short, but it is stateful, and the order of operations matters.

The failure pattern I see most often in real code reviews is not an exotic math bug. It is a copy-pasted tutorial loop with one small change in the wrong place. Someone forgets optimizer.zero_grad(). Someone validates without model.eval(). Someone logs the loss tensor itself instead of loss.item() and wonders why GPU memory keeps growing. The loop is simple. The discipline around it is what separates a model that trains cleanly from one that burns two days of GPU time to produce nonsense.

What Is Training Loop in PyTorch Explained and Why Does It Exist?

The PyTorch training loop exists because optimization is stateful, and PyTorch chooses to make that state visible rather than hide it behind a one-line fit call. Each batch passes through the same sequence: clear old gradients, run the forward pass, compute the loss, backpropagate, and step the optimizer. That looks repetitive because it is repetitive. Training is controlled repetition.

The reason this pattern matters is that gradients in PyTorch accumulate by default. Parameters remember their previous .grad values until you clear them. Autograd also records the operations from the forward pass so backward can traverse that graph in reverse. The loop is not just procedural boilerplate — it is how you manage that state correctly.

In 2026, the canonical loop usually includes a few production-grade upgrades even for ordinary models: optimizer.zero_grad(set_to_none=True) for slightly lower memory traffic, mixed precision with torch.autocast on supported GPUs, gradient clipping when the model is deep or unstable, and explicit validation blocks with model.eval() plus torch.no_grad(). If your model is compile-friendly, torch.compile can sit on top of the same loop structure without changing the fundamentals.

What does not change is the contract. The loop still answers the same four questions every iteration: what did the model predict, how wrong was it, how should the weights change, and did they actually change.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader


def train_one_epoch(model, loader, criterion, optimizer, device, scaler=None, max_grad_norm=1.0):
    model.train()
    running_loss = 0.0
    total_samples = 0

    use_amp = device.type == 'cuda'

    for inputs, labels in loader:
        inputs = inputs.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        # 1) Clear stale gradients from the previous iteration
        optimizer.zero_grad(set_to_none=True)

        # 2) Forward pass + loss calculation
        with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp):
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        # 3) Backward pass
        if scaler is not None and use_amp:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)  # unscale before clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)

            # 4) Optimizer step
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)
            optimizer.step()

        batch_size = inputs.size(0)
        running_loss += loss.item() * batch_size
        total_samples += batch_size

    return running_loss / total_samples


@torch.no_grad()
def validate(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    for inputs, labels in loader:
        inputs = inputs.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        outputs = model(inputs)
        loss = criterion(outputs, labels)

        running_loss += loss.item() * inputs.size(0)
        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)

    avg_loss = running_loss / total
    accuracy = correct / total
    return avg_loss, accuracy


# Example wiring
# model = MyClassifier().to(device)
# model = torch.compile(model)  # optional in PyTorch 2.x when the model is compile-friendly
# optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
# criterion = nn.CrossEntropyLoss()
# scaler = torch.amp.GradScaler('cuda', enabled=(device.type == 'cuda'))
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
#
# for epoch in range(10):
#     train_loss = train_one_epoch(model, train_loader, criterion, optimizer, device, scaler)
#     val_loss, val_acc = validate(model, val_loader, criterion, device)
#     scheduler.step()  # epoch-level scheduler: step after the epoch
#     lr = optimizer.param_groups[0]['lr']
#     print(f'Epoch {epoch + 1:02d} | train_loss={train_loss:.4f} | val_loss={val_loss:.4f} | val_acc={val_acc:.2%} | lr={lr:.2e}')

Experiment Tracking: Logging Metrics and Checkpoints Like an Engineer

A training loop that only prints to stdout is fine for a notebook. It is not enough for a team. Once a model matters, you need a trace of what happened: which code ran, which hyperparameters were used, what the learning rate was at each epoch, what checkpoint corresponded to the best validation metric, and when the run started to go sideways if it did.

The simple pattern is still the right one. Log epoch-level metrics to a durable store. Keep checkpoints in object storage or a mounted artifact directory. Store the checkpoint path alongside the metrics row rather than pretending those two systems are unrelated. They are part of the same training story.

The practical benefit is not just reporting. It is rollback and diagnosis. When somebody says the new model is worse, you should be able to answer with evidence: which run, which epoch, which checkpoint, what the validation curve looked like, and whether the learning rate schedule or data version changed. Without that, model debugging turns into folklore.

-- io.thecodeforge: epoch-level training metrics for experiment tracking
INSERT INTO io.thecodeforge.training_history (
    run_id,
    model_id,
    epoch_number,
    train_loss,
    val_loss,
    val_accuracy,
    learning_rate,
    checkpoint_path,
    created_at
) VALUES (
    'run_2026_04_19_001',
    'ForgeResNet50-v3',
    12,
    0.2841,
    0.3198,
    0.9142,
    0.000300,
    's3://forge-models/ForgeResNet50-v3/epoch_12.pt',
    CURRENT_TIMESTAMP
);

-- Example rollback query: fetch the best validation checkpoint for a run
SELECT epoch_number, checkpoint_path, val_loss, val_accuracy
FROM io.thecodeforge.training_history
WHERE run_id = 'run_2026_04_19_001'
ORDER BY val_accuracy DESC, val_loss ASC
LIMIT 1;

Containerizing the Forge Training Environment

Training jobs fail for boring reasons far more often than people admit: wrong CUDA runtime, mismatched drivers, DataLoader workers starved by tiny shared memory, and buffered logs that make a crashed container look idle for ten minutes. Docker does not solve those problems automatically, but it gives you one place to make them explicit.

For 2026-era PyTorch stacks, three things matter immediately. First, pin the framework and CUDA versions rather than using latest. Second, use unbuffered Python output so logs appear in real time in whatever runtime you use. Third, remember that DataLoader workers share memory through /dev/shm inside the container. If you spin up multiple workers without enough shared memory, you get hangs, worker exits, or mysterious throughput collapse.

The other trap is silent CPU fallback. Teams assume the container is using the GPU because the base image has CUDA in its name. That proves nothing. The host still needs the NVIDIA Container Toolkit, the container still needs to run with --gpus all, and your startup logs should still print torch.cuda.is_available() plus the device name. If you do not verify that, you can burn hours training on CPU and only discover it when the epoch time looks absurd.

# io.thecodeforge: Reproducible PyTorch training container
# Pin versions. Never use a floating 'latest' tag in training infrastructure.
FROM pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime

WORKDIR /app

ENV PYTHONUNBUFFERED=1
ENV PIP_NO_CACHE_DIR=1

# System deps commonly needed by vision and tabular training stacks
RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    curl \
    libgl1 \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip install -r requirements.txt

COPY . .

# Good startup hygiene: print environment info before training begins
CMD ["python", "-u", "train.py"]

# Example runtime command:
# docker run --gpus all --shm-size=2g -v $(pwd)/data:/app/data -v $(pwd)/checkpoints:/app/checkpoints forge-trainer:latest

Common Mistakes and How to Avoid Them

Most training-loop bugs come from state you forgot was stateful. Gradients persist until you clear them. Dropout and BatchNorm switch behavior based on mode. Loss tensors keep their graph unless you detach or convert them properly for logging. PyTorch is explicit about all of this, but explicit does not mean self-correcting.

The two mode switches people most often misuse are model.train() and model.eval(). These are not decorative. They change the behavior of real layers. Validation without model.eval() is not a small mistake; it changes the model you think you are measuring. The same goes for validation without torch.no_grad() — maybe the metrics are still numerically correct, but you pay the full memory cost of graphs you will never use.

The other class of mistake is device mismatch. PyTorch will never silently move your batch to match the model. If the model is on GPU and the inputs are on CPU, the first forward pass fails. That is good. What is less obvious is partial mismatch inside more complex training code — an auxiliary tensor created on CPU in the middle of loss calculation, or class weights left on CPU while logits are on GPU. The discipline is the same: decide the device once, move everything that participates in the computation there, and verify it early.

import torch


def validate_model(model, data_loader, criterion):
    device = next(model.parameters()).device
    model.eval()  # CRITICAL: switch Dropout / BatchNorm to inference behavior

    total_loss = 0.0
    total_correct = 0
    total_samples = 0

    with torch.no_grad():
        for inputs, labels in data_loader:
            inputs = inputs.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)

            outputs = model(inputs)
            loss = criterion(outputs, labels)

            total_loss += loss.item() * inputs.size(0)
            preds = outputs.argmax(dim=1)
            total_correct += (preds == labels).sum().item()
            total_samples += labels.size(0)

    avg_loss = total_loss / total_samples
    accuracy = total_correct / total_samples

    model.train()  # restore training mode for the next epoch
    return avg_loss, accuracy


# Common anti-patterns to avoid:
# 1) Forgetting model.eval() before validation
# 2) Forgetting torch.no_grad() in validation
# 3) Logging loss instead of loss.item()
# 4) Using outputs.data instead of outputs.argmax(dim=1)
# 5) Forgetting to switch back to model.train() before the next epoch

Why Your Loss Is Flat: Debugging the Silent Failures in Gradient Flow

You've seen it happen. The metrics dashboard shows a flat line for 10 epochs. Loss refuses to drop. Your first instinct might be learning rate, but seasoned engineers know the real culprit often lives deeper: vanishing or exploding gradients.

PyTorch's autograd builds a dynamic computation graph per forward pass. After backward(), gradients accumulate in .grad attributes. If you're not zeroing them with optimizer.zero_grad() each iteration, your model steps in a random direction every epoch. Worse, if your layers have sigmoid activations with inputs outside [-2,2], gradients saturate to zero. The gradient flow vanishes before it reaches early layers.

Always run a gradient sanity check early: log the mean and std of gradient norms for each layer in your first epoch. If norms drop below 1e-6 for any layer, consider switching activation functions or adding batch normalization. If you see nan or explosion (norms > 100), clip gradients with torch.nn.utils.clip_grad_norm_. It buys you debugging time without a model rewrite.

// io.thecodeforge
import torch
import torch.nn as nn

def log_gradient_norm(model: nn.Module, epoch: int, step: int):
    total_norm = 0.0
    for name, param in model.named_parameters():
        if param.grad is not None:
            param_norm = param.grad.data.norm(2)
            total_norm += param_norm.item() ** 2
            print(f"Layer '{name}' grad norm: {param_norm:.6f}")
    total_norm = total_norm ** 0.5
    print(f"Epoch {epoch} step {step} - total grad norm: {total_norm:.4f}")

Containerized Reproducibility: Why Your Model Works Locally But Fails in Production

The bug report lands on your desk: "Model accuracy drops from 94% to 55% after deployment." You check the code. Same training loop. Same checkpoint. The problem isn't your model, it's the environment. A different CUDA version, a Python stdlib change, or even numpy's random seed behavior across platforms can silently corrupt your training.

That's why you containerize everything. A Dockerfile with pinned base images fixes the runtime. But more insidious is data: if you shuffle indices or split datasets without a fixed seed, each run sees different samples. PyTorch's DataLoader shuffling depends on the system's random state. Set torch.manual_seed(42) AND DataLoader(shuffle=True, generator=torch.Generator().manual_seed(42)). Also pin torch.backends.cudnn.deterministic = True to lock cuDNN algorithms and torch.backends.cudnn.benchmark = False to avoid algorithm swapping.

Your production pipeline should match training precisely: same batch order, same augmentation transforms evaluated in the same order, same loss scaling. Break this chain, and you're debugging phantom regressions.

// io.thecodeforge
import torch
from torch.utils.data import DataLoader, TensorDataset

def build_reproducible_loader(
    data: torch.Tensor,
    labels: torch.Tensor,
    batch_size: int = 32,
    shuffle: bool = True,
    seed: int = 42
) -> DataLoader:
    torch.manual_seed(seed)
    generator = torch.Generator()
    generator.manual_seed(seed)
    dataset = TensorDataset(data, labels)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        generator=generator,
        num_workers=0  # Pin to main process for deterministic order
    )