PyTorch Training Loops

Your training run reports steadily dropping loss, but its incident router behaves erratically in staging. The loop can finish without an exception while a state bug slips by: gradients carry over between mini-batches, Dropout stays active during validation, or a checkpoint still points to live tensors that later epochs change. We'll build a small ticket router and follow each state change so you can catch those problems before they reach deployment.

We'll start after an upstream text pipeline has converted each customer ticket into two normalized features: an urgency score and a severity-language score. A binary target says whether the ticket needs specialist escalation (1) or standard routing (0). Random initial weights haven't learned to route yet, though their accuracy needn't be exactly 50%. A training loop passes examples forward, measures loss, differentiates it with respect to parameters, and applies an optimizer update. Plain SGD moves opposite the gradient; momentum and Adam transform that update.

The causal reasoning lesson uses an A/B test to ask whether a deployed update helps users. That test doesn't write the weights. The loop has to learn them first.

PyTorch doesn't hide this learning cycle behind a monolithic .fit() call: you write the loop yourself. That imperative design makes every tensor, gradient, and memory buffer inspectable at runtime. Later, when you scale up to fine-tune a large language model or run distributed training across multiple nodes, the fundamental handoffs remain identical: batches, logits, loss, gradients, optimizer updates, validation passes, and durable checkpoints.^{[1]Reference 1PyTorch: An Imperative Style, High-Performance Deep Learning Library.https://arxiv.org/abs/1912.01703[2]Reference 2Optimizing Model Parameters.https://docs.pytorch.org/tutorials/beginner/basics/optimization_tutorial.html}

The gradient descent and cross-entropy concepts build on backpropagation and softmax classification. Here we implement them in executable PyTorch code. Each code example runs self-contained on CPU unless explicitly noted otherwise. The small numbers are designed to expose tensor mechanics clearly rather than optimize classification accuracy.

These examples and all 120 plotted loss values were rechecked on PyTorch 2.14.1 on CPU on October 10, 2026. The AMP example was also executed on an RTX 5070 Ti with the CUDA-enabled 2.14.1+cu130 build. Seeds control these fixtures' random draws, but PyTorch doesn't promise identical results across versions, devices, or platforms. Compare stated properties such as a lower final loss and matching reloaded logits. Loss needn't fall on every update.^{[3]Reference 3Reproducibilityhttps://docs.pytorch.org/docs/stable/notes/randomness.html}

Turn a ticket into tensors

We'll bypass raw text tokenization for now. The feature extraction pipeline has converted each incident ticket into two normalized floating-point numbers:

Every training instance consists of a two-element feature vector and an integer class target. A mini-batch stacks those rows into a two-dimensional tensor. Six incident tickets produce a feature tensor of shape (6, 2) and a label tensor of shape (6,).

import torch

features = torch.tensor(
    [
        [-2.0, -1.2],
        [-1.3, -0.8],
        [-0.8, -1.4],
        [0.9, 1.1],
        [1.4, 0.8],
        [1.8, 1.6],
    ],
    dtype=torch.float32,
)
labels = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.long)

print("features:", tuple(features.shape), features.dtype)
print("labels:", tuple(labels.shape), labels.dtype)
print("first ticket:", features[0].tolist(), "route:", labels[0].item())

features: (6, 2) torch.float32
labels: (6,) torch.int64
first ticket: [-2.0, -1.2000000476837158] route: 0

For this hard-label classifier, floating-point features (torch.float32) match the linear layer's weights, while torch.long targets supply class indices. Floating targets of shape (batch,) don't satisfy that contract. Cross-entropy also supports a different contract: floating soft targets of shape (batch, classes), with a probability distribution in each row. PyTorch doesn't validate those distributions for you.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html}

The shapes match the loss. The two numbers for one ticket are still not a probability that the ticket should be escalated.

Logits are scores, not probabilities

A linear classification head ending in nn.Linear(2, 2) produces two unconstrained scores for each ticket:

These unnormalized outputs are logits. They can be positive, negative, or zero and needn't sum to one. For our hard-label classifier, nn.CrossEntropyLoss combines log-softmax over raw logits with negative log-likelihood at integer target indices.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html}

Consider raw ticket logits [2.0, -1.0] for an example whose ground-truth class is 0. To see how cross-entropy scores this prediction, convert the logits to probabilities:

The second ticket outputs logits [-0.5, 1.5] with ground-truth class 1. The correct class wins, but with less margin: $p_1 \approx 0.8808$, yielding loss $-\log(p_1) \approx 0.1269$. The average loss over both tickets is $(0.0486 + 0.1269) / 2 \approx 0.0878$. PyTorch's default reduction="mean" computes exactly this batch average.

import torch
from torch import nn

logits = torch.tensor([[2.0, -1.0], [-0.5, 1.5]], dtype=torch.float32)
labels = torch.tensor([0, 1], dtype=torch.long)
loss = nn.CrossEntropyLoss()(logits, labels)

routes = ["standard", "escalate"]
predictions = logits.argmax(dim=1).tolist()
print("predicted routes:", [routes[index] for index in predictions])
print("cross entropy:", round(loss.item(), 4))

predicted routes: ['standard', 'escalate']
cross entropy: 0.0878

Failure case: Don't pass torch.softmax(logits, dim=-1) as the input to CrossEntropyLoss. It treats those probabilities as logits, giving gradients for a different objective. Separately, torch.log(torch.softmax(logits, dim=-1)) can turn an underflowed probability into -inf. Direct log-softmax/cross-entropy avoids that intermediate rounding through log-sum-exp: $\log \sum_j e^{z_j}=c+\log\sum_j e^{z_j-c}$, with $c=\max_j z_j$.

CrossEntropyLoss summarizes batch predictions into a single scalar value. That scalar says how wrong the batch is. It still hasn't moved a weight.

The five execution phases in a training step

This eager SGD loop uses five phases. Clearing gradients must precede the backward pass whose derivatives you want to retain; it needn't precede the forward pass specifically. Before checking the table, identify the phase that applies the parameter update:

In this loop, loss.backward() doesn't modify weight values. It traverses the graph through the chain rule and accumulates derivatives in reachable leaf tensors with requires_grad=True. An unused parameter needn't receive a gradient.^{[5]Reference 5Deep Learning.https://www.deeplearningbook.org/[6]Reference 6Autograd mechanicshttps://docs.pytorch.org/docs/stable/notes/autograd.html} A training forward pass can change buffers such as BatchNorm's running statistics, even though the table's parameter values remain unchanged.

The parameter values themselves remain completely unchanged until optimizer.step() runs. The optimizer reads the accumulated gradients, applies its update rule (such as SGD or Adam), and updates the parameter storage in place.

Why use set_to_none=True, the current default? It removes the parameter's reference to its gradient tensor and avoids explicitly zero-filling that tensor. This generally reduces memory use and can modestly improve performance. Other references or an allocator's cache can retain the storage, leaving operating-system GPU-memory reports unchanged.^{[7]Reference 7Optimizer.zero_gradhttps://docs.pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html}

It also changes optimization semantics. A parameter receiving no gradient remains .grad = None, and standard PyTorch optimizers skip it. An existing zero gradient still participates in an update: accumulated momentum or configured weight decay can move the parameter. Neither effect is inevitable when those terms are absent. For custom gradient inspection, check p.grad is not None before accessing its values.

Trace this step with exact arithmetic. Consider two symmetric tickets: ticket [1.0, 1.0] has label 1 (escalate), and ticket [-1.0, -1.0] has label 0 (standard). Let $W$ be a 2x2 weight matrix where row 0 scores class 0 and row 1 scores class 1. Initializing $W = 0$ with zero bias yields logits [0.0, 0.0] for both tickets. The predicted class probabilities are [0.5, 0.5], producing a cross-entropy loss of $-\log(0.5) \approx 0.6931$ for each row.

For cross-entropy loss with mean reduction, each row's gradient contribution to weight $W_{c,j}$ is $(p_c - y_c) \cdot x_j / N$, where $N=2$ is the batch size. Both tickets produce positive gradient values for class 0 weights and negative gradient values for class 1 weights, yielding an aggregate gradient matrix $dW = \begin{bmatrix} 0.5 & 0.5 \\ -0.5 & -0.5 \end{bmatrix}$. The bias gradients cancel to zero. loss.backward() leaves $W$ at zero while writing $dW$ into $W\text{.grad}$. Finally, SGD with learning rate $\eta = 0.2$ updates the weights:

View larger

The arithmetic says backward() fills the gradient and step() moves the weights. The same two tickets should show that split in PyTorch.

import torch
from torch import nn

features = torch.tensor([[1.0, 1.0], [-1.0, -1.0]], dtype=torch.float32)
labels = torch.tensor([1, 0], dtype=torch.long)
model = nn.Linear(2, 2)
with torch.no_grad():
    model.weight.zero_()
    model.bias.zero_()

optimizer = torch.optim.SGD(model.parameters(), lr=0.2)
loss_fn = nn.CrossEntropyLoss()

before = model.weight.detach().clone()
optimizer.zero_grad(set_to_none=True)
loss = loss_fn(model(features), labels)
loss.backward()
gradient = model.weight.grad.detach().clone()
after_backward = model.weight.detach().clone()
optimizer.step()
after_step = model.weight.detach().clone()

print("loss:", round(loss.item(), 4))
print("gradient filled:", bool(gradient.abs().sum() > 0))
print("changed by backward:", not torch.equal(before, after_backward))
print("changed by step:", not torch.equal(before, after_step))
torch.testing.assert_close(gradient, torch.tensor([[0.5, 0.5], [-0.5, -0.5]]))
torch.testing.assert_close(after_step, torch.tensor([[-0.1, -0.1], [0.1, 0.1]]))

loss: 0.6931
gradient filled: True
changed by backward: False
changed by step: True

The output confirms the state boundaries: loss.backward() populated .grad without touching parameter values, and optimizer.step() moved the parameters. What happens if leftover gradients from an earlier iteration remain in .grad when the next backward() executes?

Clearing gradients and micro-batch accumulation

PyTorch adds newly computed gradients into existing .grad buffers rather than overwriting them (param.grad += new_grad). That accumulation default is intentional: it enables gradient accumulation, letting multiple micro-batches build up an effective batch update when memory is constrained. In standard training where each mini-batch should trigger an independent step, failing to clear gradients causes stale derivatives to pollute future updates.

To observe this accumulation directly, differentiate $y = w^2$ twice at $w = 1.0$. The derivative is $\frac{dy}{dw} = 2w = 2.0$. Calling backward() a second time without clearing gradients adds another $2.0$, producing an accumulated gradient of $4.0$. Resetting weight.grad = None restores the clean single-step value:

import torch

weight = torch.tensor(1.0, requires_grad=True)

(weight**2).backward()
first_grad = weight.grad.item()

(weight**2).backward()
accumulated_grad = weight.grad.item()

weight.grad = None
(weight**2).backward()
cleared_grad = weight.grad.item()

print("first backward:", first_grad)
print("without clearing:", accumulated_grad)
print("after clearing:", cleared_grad)

first backward: 2.0
without clearing: 4.0
after clearing: 2.0

Without clearing, the stored derivative doubles. Clearing before an independent batch's backward pass removes stale loss gradients. An optimizer with momentum still carries information from earlier updates in its own state.^{[7]Reference 7Optimizer.zero_gradhttps://docs.pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html}

Gradient accumulation across micro-batches

When a model or batch size exceeds available GPU memory, gradient accumulation lets you simulate a larger effective batch size $B_{\text{eff}} = K \times B_{\text{micro}}$ using $K$ smaller micro-batches of size $B_{\text{micro}}$.

The loss function usually averages errors over the current micro-batch: $L_k = \frac{1}{B} \sum_{i=1}^B \ell(x_{k,i}, y_{k,i})$. The true mean loss over the combined batch of $K \cdot B$ examples is:

For equal-sized micro-batches with a per-example mean loss, (loss / K).backward() recovers the combined mean gradient. Without division, that gradient is $K$ times larger. For plain SGD without other terms, this scales the immediate update like a $K$-times learning rate; clipping, momentum history, and Adam needn't have that equivalence.^{[8]Reference 8Automatic Mixed Precision Exampleshttps://docs.pytorch.org/docs/stable/notes/amp_examples.html}

An incomplete accumulation window needs its actual denominator. With micro-batch sizes $n_k$, weight each mean by $n_k / \sum_j n_j$. Dividing two losses equally when their batches contain two rows and one row gives that last ticket too much influence. Predict which of these two accumulated gradients matches the full three-ticket batch:

import copy
import torch
from torch import nn

features = torch.tensor([[1., 1.], [-1., -1.], [2., 1.]], dtype=torch.float64)
labels = torch.tensor([1, 0, 1])
base = nn.Linear(2, 2, dtype=torch.float64)
with torch.no_grad():
    base.weight.zero_()
    base.bias.zero_()
loss_fn = nn.CrossEntropyLoss()

full = copy.deepcopy(base)
loss_fn(full(features), labels).backward()

def accumulate(row_weighted):
    model = copy.deepcopy(base)
    chunks = [(features[:2], labels[:2]), (features[2:], labels[2:])]
    for xb, yb in chunks:
        weight = len(yb) / len(labels) if row_weighted else 1 / len(chunks)
        (weight * loss_fn(model(xb), yb)).backward()
    return model.weight.grad

equal_chunks = accumulate(False)
by_rows = accumulate(True)
print("full-batch class-0 gradient:", [round(x, 3) for x in full.weight.grad[0].tolist()])
print("equal-chunk class-0 gradient:", [round(x, 3) for x in equal_chunks[0].tolist()])
print("row weighting matches:", torch.allclose(by_rows, full.weight.grad))
torch.testing.assert_close(by_rows, full.weight.grad)

full-batch class-0 gradient: [0.667, 0.5]
equal-chunk class-0 gradient: [0.75, 0.5]
row weighting matches: True

That equivalence assumes unchanged weights across the window and matching per-example computations. BatchNorm statistics depend on each micro-batch, and Dropout draws can differ, so accumulation isn't generally identical to one physical large-batch forward pass. With weighted or ignored-label losses, use the loss's effective denominator rather than raw row counts. Step once per window and clear afterward. With AMP, keep one scale throughout the window, then unscale, clip, step, and update the scale once.

DDP normally synchronizes gradient buckets during backward. For accumulation, wrap both forward and backward for the first $K-1$ micro-batches in model.no_sync(); the final backward outside that context synchronizes the accumulated gradients. This suppresses gradient synchronization, not necessarily every distributed communication. Match the accumulation schedule across ranks.^{[9]Reference 9DistributedDataParallelhttps://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html}

One update on a synthetic fixture verifies autograd mechanics, but it doesn't constitute a complete loop. We need abstractions to manage datasets and partition rows into mini-batches.

Batching and row-weighted epoch loss

A dataset exposes samples and labels; it can read or generate them lazily. A data loader coordinates batching, optional worker processes, and sampling. For independent tickets, shuffling changes the order across epochs and reduces systematic ordering effects. Gradient diversity and oscillation also depend on the data, batches, model, and optimizer. Preserve meaningful sequence order when the task requires it.

Using our six tickets with batch_size=4, the data loader yields two batches: the first contains four rows, and the second contains the remaining two rows.

import torch
from torch.utils.data import DataLoader, TensorDataset

features = torch.tensor(
    [
        [-2.0, -1.2],
        [-1.3, -0.8],
        [-0.8, -1.4],
        [0.9, 1.1],
        [1.4, 0.8],
        [1.8, 1.6],
    ],
    dtype=torch.float32,
)
labels = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.long)
loader = DataLoader(TensorDataset(features, labels), batch_size=4, shuffle=False)

for batch_index, (xb, yb) in enumerate(loader, start=1):
    print(f"batch {batch_index}: shape={tuple(xb.shape)} labels={yb.tolist()}")

batch 1: shape=(4, 2) labels=[0, 0, 0, 1]
batch 2: shape=(2, 2) labels=[1, 1]

The final batch contains only two rows. When calculating average loss across an epoch, don't simply average the reported batch loss numbers. Because batch sizes differ, taking an unweighted mean of batch averages assigns excessive influence to the smaller batch.

For this unweighted loss without ignored labels, multiply each batch mean by its row count and divide the sum by rows actually processed. Batches of size 4 and 2 with losses 0.20 and 0.50 yield $(4 \times 0.20 + 2 \times 0.50) / 6 = 0.30$, versus the incorrect equal-batch mean 0.35. Class weights and ignored labels can change the denominator.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html} During training, this summarizes losses at changing weights, not a fresh evaluation of the final checkpoint.

With data loading established, we can construct a minimal training loop. Before evaluating generalization on held-out data, we test whether the model can memorize a small, clean batch.

Overfitting one batch as an isolation test

Before scaling up, run an overfit-one-batch test on a clean fixture the model has enough capacity to fit. Temporarily disable regularization that would obstruct memorization. If loss remains high, inspect labels, shapes, the loss contract, learning rate, and update ordering. Conflicting labels for identical inputs or insufficient capacity can prevent near-zero loss even with a correct loop.

We'll define a multi-layer classifier: nn.Linear(2, 8) expands the two inputs to eight hidden activations, nn.ReLU() introduces non-linearity by zeroing negative values, and nn.Linear(8, 2) produces class logits.

import torch
from torch import nn

torch.manual_seed(7)
features = torch.tensor(
    [
        [-2.0, -1.2], [-1.3, -0.8], [-0.8, -1.4], [-1.6, -0.4],
        [0.9, 1.1], [1.4, 0.8], [1.8, 1.6], [0.7, 1.7],
    ],
    dtype=torch.float32,
)
labels = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.long)
model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 2))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.15)

with torch.no_grad():
    initial_loss = loss_fn(model(features), labels).item()

for _ in range(120):
    optimizer.zero_grad(set_to_none=True)
    loss = loss_fn(model(features), labels)
    loss.backward()
    optimizer.step()

with torch.no_grad():
    final_logits = model(features)
    final_loss = loss_fn(final_logits, labels).item()
    accuracy = (final_logits.argmax(dim=1) == labels).float().mean().item()

print("initial loss:", round(initial_loss, 4))
print("final loss:", round(final_loss, 4))
print("memorized batch:", accuracy == 1.0)

initial loss: 0.6521
final loss: 0.0056
memorized batch: True

The loss drops from 0.6521 to 0.0056 and classification accuracy reaches 100%. This passes the basic wiring check. Now we can separate training from held-out validation.

Isolating validation from training updates

The loss reduction checks learning on this fixture; memorization alone doesn't establish generalization to unseen tickets. Split the workflow into two phases:

• Training phase: Mini-batches flow forward, autograd constructs the computation graph, loss.backward() calculates derivatives, and optimizer.step() updates parameters.
• Validation phase: Held-out examples evaluate the current weights without updating parameters or recording gradient history.

Validation relies on two independent controls:

The script below trains on eight tickets and tracks validation loss across four held-out tickets over 60 epochs.

import copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(7)
train_x = torch.tensor(
    [
        [-2.0, -1.2], [-1.3, -0.8], [-0.8, -1.4], [-1.6, -0.4],
        [0.9, 1.1], [1.4, 0.8], [1.8, 1.6], [0.7, 1.7],
    ],
    dtype=torch.float32,
)
train_y = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.long)
val_x = torch.tensor([[-1.1, -0.6], [-0.5, -1.7], [1.0, 0.6], [1.7, 0.4]])
val_y = torch.tensor([0, 0, 1, 1], dtype=torch.long)

loader = DataLoader(TensorDataset(train_x, train_y), batch_size=4, shuffle=True)
model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 2))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.15)
best_loss = float("inf")
best_state = None
best_epoch = None

for epoch in range(1, 61):
    model.train()
    train_total = 0.0
    train_count = 0
    for xb, yb in loader:
        optimizer.zero_grad(set_to_none=True)
        loss = loss_fn(model(xb), yb)
        if not torch.isfinite(loss):
            raise RuntimeError(f"non-finite training loss at epoch {epoch}")
        loss.backward()
        optimizer.step()
        train_total += loss.item() * len(yb)
        train_count += len(yb)

    model.eval()
    with torch.no_grad():
        val_logits = model(val_x)
        val_loss = loss_fn(val_logits, val_y).item()
    if not torch.isfinite(torch.tensor(val_loss)):
        raise RuntimeError(f"non-finite validation loss at epoch {epoch}")

    if val_loss < best_loss:
        best_loss = val_loss
        best_epoch = epoch
        best_state = copy.deepcopy(model.state_dict())

assert best_state is not None
model.load_state_dict(best_state)
model.eval()
with torch.no_grad():
    selected_accuracy = (model(val_x).argmax(dim=1) == val_y).float().mean().item()

print("best validation loss:", round(best_loss, 4))
print("selected epoch:", best_epoch)
print("selected validation accuracy:", round(selected_accuracy, 3))
print("saved tensors:", len(best_state))

best validation loss: 0.0118
selected epoch: 60
selected validation accuracy: 1.0
saved tensors: 4

The checkpoint snapshot aliasing trap

Notice copy.deepcopy(model.state_dict()) when validation improves. A plain best_state = model.state_dict() is a shallow snapshot whose tensors share live parameter and persistent-buffer storage. Subsequent in-place updates change those values. Deep-copy the state or serialize it immediately to preserve the selected checkpoint.^{[10]Reference 10Saving and Loading Modelshttps://docs.pytorch.org/tutorials/beginner/saving_loading_models.html}

We can demonstrate this reference leak directly:

import copy
import torch
from torch import nn

model = nn.Linear(2, 2)
with torch.no_grad():
    model.weight.zero_()

aliased = model.state_dict()
snapshot = copy.deepcopy(model.state_dict())
with torch.no_grad():
    model.weight.add_(1.0)

print("plain state followed the update:", torch.equal(aliased["weight"], model.weight))
print("copied state stayed at zero:", torch.count_nonzero(snapshot["weight"]).item() == 0)

plain state followed the update: True
copied state stayed at zero: True

View larger

These four validation tickets select the checkpoint, so their reported 100% accuracy isn't an independent test of that selection. Use an untouched test set for the final estimate and a representative deployment evaluation. The loop calls model.eval() even without mode-sensitive layers; adding Dropout makes its purpose visible.

Operational modes and context managers

In a network composed strictly of linear transformations and ReLU activations, model.train() and model.eval() yield identical outputs because neither layer maintains mode-dependent state. In production architectures, however, several layers change their behavior across modes:

• Dropout: For $0 \leq p < 1$, training zeroes activations randomly and scales survivors by $1/(1-p)$. Evaluation is identity. At $p=1$, training zeroes everything.
• BatchNorm: With default track_running_stats=True, training uses batch statistics and updates running buffers; evaluation uses the running statistics. With track_running_stats=False, it uses batch statistics in both modes. A cumulative average is also possible with momentum=None.^{[11]Reference 11BatchNorm2dhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.BatchNorm2d.html}
• LayerNorm: It normalizes each input over the trailing dimensions specified by normalized_shape, often hidden features. The layer itself has the same computation in both modes.

The snippet below introduces Dropout to our classifier, runs the same input ticket twice in training mode, and then runs it twice in evaluation mode.

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(2, 6), nn.ReLU(), nn.Dropout(p=0.5), nn.Linear(6, 2))
ticket = torch.tensor([[1.2, 0.9]])

model.train()
train_a = model(ticket)
train_b = model(ticket)

model.eval()
with torch.no_grad():
    eval_a = model(ticket)
    eval_b = model(ticket)

print("training outputs equal:", torch.equal(train_a, train_b))
print("evaluation outputs equal:", torch.equal(eval_a, eval_b))

training outputs equal: False
evaluation outputs equal: True

This seeded pair differs in training mode. Two Dropout calls can sample the same mask and produce identical outputs. model.eval() controls layer behavior, while torch.no_grad() controls graph recording.

Saving weights to disk produces the deployable artifact. Checkpoints also protect ongoing training against sudden hardware faults and preemption.

Resilient checkpointing and atomic state snapshots

Save a reloadable artifact with its architecture, feature schema, and label mapping. Prefer a state_dict over torch.save(model, path) for this workflow. Whole-module pickling records references to Python classes and import paths; it doesn't package their class definitions, and refactoring those paths can break loading.^{[10]Reference 10Saving and Loading Modelshttps://docs.pytorch.org/tutorials/beginner/saving_loading_models.html}

The model's state dictionary normally contains named parameter and persistent-buffer tensors; custom modules can also supply extra state. It doesn't automatically capture architecture code, ordinary Python attributes, or accumulated gradients.

The complete resumption bundle

For resumption, identify the state your particular loop needs. At an update/epoch boundary, a typical bundle includes:

1. model_state_dict: Trainable parameters and persistent buffers (such as BatchNorm running stats).
2. optimizer_state_dict: Adam moments and counters, or SGD momentum if configured. Omitting them changes subsequent updates; degraded convergence or oscillation isn't inevitable.^{[10]Reference 10Saving and Loading Modelshttps://docs.pytorch.org/tutorials/beginner/saving_loading_models.html}
3. lr_scheduler_state_dict: Current learning rate step counter and cycle phase.
4. scaler_state_dict: The dynamic loss scale factor and consecutive finite step count when using Automatic Mixed Precision.
5. rng_state: Current states of every RNG used, including PyTorch CPU/CUDA, Python, NumPy, and explicit DataLoader generators. A seed restarts a sequence; it doesn't restore the current position.^{[3]Reference 3Reproducibilityhttps://docs.pytorch.org/docs/stable/notes/randomness.html}
6. Metadata: The epoch index, global step count, validation metrics, git commit hash, and dataset split identifier.

Mid-epoch resumption can also require sampler position, worker/stream state, and buffered data. Mid-window gradient accumulation needs accumulated gradients or a replay of the unfinished window; state_dict() doesn't include parameter .grad. Even with these states, bit-for-bit continuation depends on compatible software, hardware, and deterministic operations. Saving at a clear update boundary makes recovery easier to reason about.

Preventing corruption with atomic file writes

If a process is interrupted while overwriting its only checkpoint, that file may be truncated or incomplete. Reloading can raise an EOF, ZIP/archive, or other serialization error. The exact message depends on the format and damaged bytes. Keeping older checkpoints limits how much progress one failed save can lose.

To make checkpointing resilient against preemption, write atomically:

1. Serialize to a staging file in the destination directory.
2. Flush Python buffers and os.fsync the staging file. Check for failures.
3. Successfully replace the destination with os.replace. POSIX requires an atomic rename, so readers see the old file or the completed replacement. If there was no previous checkpoint, interruption before replacement can leave no final file.
4. On Linux, sync the parent directory too so the renamed directory entry reaches storage.^{[12]Reference 12os: fsync and replacehttps://docs.python.org/3.14/library/os.html[13]Reference 13fsync(2): file and directory synchronizationhttps://man7.org/linux/man-pages/man2/fsync.2.html}

Atomic visibility and power-loss durability are separate properties. This pattern protects a previous valid file against interruption during staging; it doesn't promise immunity to storage faults or every network/object-storage contract. Coordinate writers and use unique staging names when saves can overlap.

View larger

Since PyTorch 2.6, torch.load defaults to weights_only=True when no pickle_module is supplied. Its restricted unpickler narrows the code-execution attack surface; it isn't a general sandbox or protection against every malicious file. Some full resumption bundles, such as raw NumPy RNG objects, need compatible encoding or trusted, deliberate allowlisting. Keep the explicit flag here.^{[14]Reference 14Serialization Semanticshttps://docs.pytorch.org/docs/stable/notes/serialization.html}

The next example has one writer and a local file. It reloads both weights and SGD momentum, then checks the next update, including a weights-only copy that deliberately omits momentum. It also raises a controlled exception during a second staging write and verifies that the live checkpoint stays intact. This tests the staging failure path, not a power loss or filesystem fault. Its directory-sync branch targets Linux.

import os
import sys
from pathlib import Path
from tempfile import TemporaryDirectory

import torch
from torch import nn

torch.manual_seed(11)
features = torch.tensor([[-1.5, -1.0], [-0.9, -1.2], [1.1, 0.8], [1.6, 1.4]])
labels = torch.tensor([0, 0, 1, 1], dtype=torch.long)
model = nn.Linear(2, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.2, momentum=0.9)
loss_fn = nn.CrossEntropyLoss()

for _ in range(80):
    optimizer.zero_grad(set_to_none=True)
    loss = loss_fn(model(features), labels)
    loss.backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    original_logits = model(features)
    original = original_logits.argmax(dim=1)

with TemporaryDirectory() as directory:
    final_path = Path(directory) / "incident-router.pt"
    tmp_path = Path(directory) / "incident-router.pt.tmp"

    checkpoint = {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "model_config": {"in_features": 2, "out_features": 2},
        "feature_names": ["urgency", "severity_language"],
        "label_names": ["standard", "escalate"],
        "training_steps": 80,
    }

    with open(tmp_path, "wb") as f:
        torch.save(checkpoint, f)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp_path, final_path)
    if sys.platform.startswith("linux"):
        directory_fd = os.open(final_path.parent, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(directory_fd)
        finally:
            os.close(directory_fd)

    previous_bytes = final_path.read_bytes()
    try:
        with open(tmp_path, "wb") as f:
            f.write(b"incomplete replacement")
            raise RuntimeError("simulated interruption before replace")
    except RuntimeError as error:
        assert str(error) == "simulated interruption before replace"
    staging_failure_preserved_live = final_path.read_bytes() == previous_bytes
    assert staging_failure_preserved_live

    loaded = torch.load(final_path, map_location="cpu", weights_only=True)
    restored = nn.Linear(**loaded["model_config"])
    restored.load_state_dict(loaded["model_state_dict"])
    restored_optimizer = torch.optim.SGD(restored.parameters(), lr=0.2, momentum=0.9)
    restored_optimizer.load_state_dict(loaded["optimizer_state_dict"])
    weights_only = nn.Linear(**loaded["model_config"])
    weights_only.load_state_dict(loaded["model_state_dict"])
    fresh_optimizer = torch.optim.SGD(weights_only.parameters(), lr=0.2, momentum=0.9)
    restored.eval()
    with torch.no_grad():
        reloaded_logits = restored(features)
        reloaded = reloaded_logits.argmax(dim=1)

torch.testing.assert_close(original_logits, reloaded_logits)
for candidate, candidate_optimizer in (
    (model, optimizer), (restored, restored_optimizer), (weights_only, fresh_optimizer)
):
    candidate.train()
    candidate_optimizer.zero_grad(set_to_none=True)
    loss_fn(candidate(features), labels).backward()
    candidate_optimizer.step()
for parameter, restored_parameter in zip(model.parameters(), restored.parameters()):
    torch.testing.assert_close(parameter, restored_parameter)

print("routes:", original.tolist())
print("reload agrees:", torch.equal(original, reloaded))
print("labels:", loaded["label_names"])
print("has optimizer state:", "optimizer_state_dict" in loaded)
print("next update agrees:", all(torch.equal(a, b) for a, b in zip(model.parameters(), restored.parameters())))
print("weights-only next update agrees:", all(torch.equal(a, b) for a, b in zip(model.parameters(), weights_only.parameters())))
print("staging failure preserved live file:", staging_failure_preserved_live)

routes: [0, 0, 1, 1]
reload agrees: True
labels: ['standard', 'escalate']
has optimizer state: True
next update agrees: True
weights-only next update agrees: False
staging failure preserved live file: True

The reloaded logits match, and the fully restored copy takes the same next momentum update on a fixed batch. The weights-only copy starts at those same weights but takes a different update because its momentum history is missing. The incomplete staging file never replaces the valid live file. These checks don't test shuffled sampling, Dropout RNG restoration, a scheduler, process termination, or power-loss recovery.

Diagnosing failures and bounding gradients

When a training script runs without throwing exceptions, subtle bugs can still compromise learning. Use this triage matrix to isolate root causes:

Ordinary backward releases its saved tensors; appending a loss afterward doesn't necessarily retain all its forward activations. It still keeps the loss tensor and graph references alive. Detaching removes graph history but leaves a tensor on its current device. For scalar reporting, .item() produces a Python number; on CUDA, that read can wait for GPU work, so choose a reporting cadence separately from the update loop.^{[6]Reference 6Autograd mechanicshttps://docs.pytorch.org/docs/stable/notes/autograd.html[15]Reference 15CUDA semanticshttps://docs.pytorch.org/docs/2.14/notes/cuda.html}

Five-step numerical triage when loss becomes NaN

When loss becomes non-finite, don't immediately reduce the learning rate blindly. Follow this ordered diagnostic sequence to find the earliest invalid value:

1. Inputs: Check raw batch inputs (torch.isfinite(xb).all()). Unsanitized nulls or extreme feature values corrupt the forward pass immediately.
2. Logits: Check model output scores (torch.isfinite(logits).all()). If inputs are finite but logits are NaN, examine numerical stability in custom layers or activations.
3. Loss: Check torch.isfinite(loss), then target values, weights, masks, and reductions. All-ignored targets can give an undefined mean; finite extreme logits can overflow a loss difference. Out-of-range hard targets ordinarily raise an indexing error instead of explaining every Inf.
4. Gradients: After backward (and AMP unscaling), check torch.isfinite(p.grad).all() only where p.grad is not None. Loss and gradient finiteness are separate checks; inspect both.
5. Parameters: After optimizer.step(), check updated weights (torch.isfinite(p).all()).

Bounding gradients with global norm clipping

When dealing with deep recurrent networks or transformer architectures, gradient vectors can intermittently surge in magnitude. Gradient clipping enforces an upper bound on gradient size before the optimizer applies parameter updates.

PyTorch's nn.utils.clip_grad_norm_ calculates the total $\ell_2$ norm across all model parameters concatenated into a single vector:

If $g_{\text{norm}}$ exceeds max_norm, the function scales all gradients down proportionally:

For a finite nonzero gradient and positive bound, global norm clipping rescales all components together, preserving their direction up to rounding. The expression above is the ideal rule; PyTorch's implementation uses a small denominator epsilon. Value clipping clamps individual entries and can change direction.^{[16]Reference 16clip_grad_norm_https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html} Clipping bounds the gradient, not every optimizer's parameter displacement: momentum, Adam, and weight decay can transform the update.

import torch
from torch import nn

parameter = nn.Parameter(torch.tensor([3.0, 4.0]))
parameter.grad = torch.tensor([6.0, 8.0])

before = torch.linalg.vector_norm(parameter.grad).item()
reported_norm = nn.utils.clip_grad_norm_([parameter], max_norm=5.0, error_if_nonfinite=True).item()
after = torch.linalg.vector_norm(parameter.grad).item()

print("norm before:", round(before, 1))
print("reported norm:", round(reported_norm, 1))
print("norm after:", round(after, 1))
print("finite:", torch.isfinite(parameter.grad).all().item())

norm before: 10.0
reported norm: 10.0
norm after: 5.0
finite: True

error_if_nonfinite=True raises when the computed total norm is non-finite, including overflow while combining otherwise finite values. Check the optimizer update separately to confirm that it remains finite.

Inference mode versus no-grad contexts

During deployment or batch inference, the model generates predictions without ever calculating gradients. Autograd tracking is entirely unnecessary.

PyTorch offers two context managers for non-gradient execution:

• torch.no_grad() suppresses reverse-mode graph recording while retaining ordinary tensor tracking. Its outputs can later be used as constants in a recorded computation.
• torch.inference_mode() removes additional view/version tracking for inference tensors and disables forward-mode AD too. Those tensors are restricted when a later recorded computation needs to save them for backward. Use it for isolated inference; use no-grad when outputs feed training. Benchmark any speed gain rather than assuming one.^{[17]Reference 17inference_modehttps://docs.pytorch.org/docs/2.14/generated/torch.autograd.grad_mode.inference_mode.html}

Neither context switches the model to evaluation mode. For example, extracting frozen features with inference mode and multiplying them by trainable weights can raise Inference tensors cannot be saved for backward. Extract with no-grad, or clone the inference tensor outside that context to create an ordinary tensor.

import torch
from torch import nn

model = nn.Linear(2, 2)
with torch.no_grad():
    model.weight.copy_(torch.tensor([[-0.1, -0.1], [0.1, 0.1]]))
    model.bias.zero_()
model.eval()
ticket = torch.tensor([[1.5, 1.2]])

with torch.inference_mode():
    logits = model(ticket)
    prediction = logits.argmax(dim=1).item()

print("gradient tracking:", logits.requires_grad)
print("logits:", [round(value, 2) for value in logits[0].tolist()])
print("route class:", prediction)

head = nn.Linear(2, 2)
try:
    head(logits).sum().backward()
except RuntimeError as error:
    if "Inference tensors cannot be saved for backward" not in str(error):
        raise
    print("inference features rejected by trainable head:", True)
else:
    raise AssertionError("expected the head to reject inference tensors")

ordinary_features = logits.clone()  # Outside inference_mode.
head(ordinary_features).sum().backward()
assert head.weight.grad is not None and torch.isfinite(head.weight.grad).all()
print("cloned features permit head backward:", True)

gradient tracking: False
logits: [-0.27, 0.27]
route class: 1
inference features rejected by trainable head: True
cloned features permit head backward: True

The prediction path needs no gradients, but the separate trainable head does. The clone restores an ordinary tensor without recreating a gradient path through the frozen router. Using no_grad() for extraction would also produce ordinary features without this repair. The next boundary changes numerical precision during training.

Automatic mixed precision and dynamic loss scaling

Our eager examples use float32 parameters and arithmetic. Compatible GPUs can run suitable lower-precision matrix operations faster. A float16/BF16 activation tensor has half the float32 payload, but total training memory also includes parameters, gradients, optimizer state, and other tensors; AMP doesn't halve all of it.

Automatic Mixed Precision (AMP) uses device-specific operation policies. CUDA autocast can run linear/matrix operations in float16 or BF16 and operations such as cross-entropy in float32. Unlisted operations, explicit dtypes, and in-place calls have different rules. Keep the model's stored parameters in float32 for this pattern; autocast supplies lower-precision inputs to eligible operations.^{[18]Reference 18Mixed Precision Training.https://arxiv.org/abs/1710.03740[19]Reference 19Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

The float16 underflow hazard and GradScaler

Float16 has 5 exponent bits and 10 stored fraction bits. Its smallest positive normal value is $2^{-14} \approx 6.1 \times 10^{-5}$, but subnormals extend to $2^{-24} \approx 5.96 \times 10^{-8}$. A value such as $10^{-5}$ isn't intrinsically zero in the format. Smaller values can round to zero, and some execution paths flush subnormals. Tiny intermediate backward gradients can therefore lose information even when the final parameter buffers are float32.^{[20]Reference 20Type Infohttps://docs.pytorch.org/docs/stable/type_info.html}

torch.amp.GradScaler reduces this risk through dynamic loss scaling:

1. Forward pass: Operations execute inside torch.amp.autocast('cuda'). CrossEntropyLoss runs in float32 for numerical stability.
2. Loss scaling: The initial default scale is $S=65,536$. Holding it constant during differentiation gives $\nabla_W(SL)=S\nabla_W L$. Scaling raises small backward values, but one scale can't guarantee every intermediate fits the format.
3. Backward pass: Calling scaler.scale(loss).backward() calculates scaled gradients ($S \cdot \nabla W$).
4. Unscaling before clipping: Call scaler.unscale_(optimizer) once, after accumulation and before clipping. Otherwise the threshold applies to scaled gradients. With $S=65,536$ and max_norm=1, clipping first imposes an unscaled bound of roughly $1/65,536$ instead of 1.
5. Conditional optimizer step: For this standard optimizer, scaler.step calls its step only if the non-finite check passed. Otherwise it skips the update. It doesn't clear .grad; the next zero_grad does that. Finite gradients still don't prove a sound optimizer update.
6. Dynamic scale update: With defaults, a non-finite event halves $S$ and resets the growth counter. After 2,000 consecutive successful iterations, $S$ doubles. Ordinary finite steps leave it unchanged; these controls are configurable. The scale can fall below 1 when float16 overflow is the problem.^{[21]Reference 21GradScaler implementation (PyTorch 2.14.0)https://github.com/pytorch/pytorch/blob/v2.14.0/torch/amp/grad_scaler.py}

View larger

BF16 has the same exponent width as float32, making float16-style range problems much less likely. It still has finite range and less significand precision; it doesn't eliminate underflow, overflow, or rounding. BF16 autocast commonly runs without loss scaling. That differs from manually casting every model parameter to BF16.

import torch
from torch import nn

torch.manual_seed(7)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
use_amp = device.type == "cuda"
model = nn.Linear(2, 2).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.CrossEntropyLoss()
xb = torch.tensor([[1.2, 0.8], [-1.0, -0.7]], device=device)
yb = torch.tensor([1, 0], dtype=torch.long, device=device)
scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
before = [parameter.detach().clone() for parameter in model.parameters()]

optimizer.zero_grad(set_to_none=True)
with torch.amp.autocast("cuda", dtype=torch.float16, enabled=use_amp):
    logits = model(xb)
    loss = loss_fn(logits, yb)

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

changed = any(not torch.equal(old, new) for old, new in zip(before, model.parameters()))
print("AMP enabled:", use_amp)
print("finite loss:", bool(torch.isfinite(loss)))
print("finite unscaled gradient norm:", bool(torch.isfinite(grad_norm)))
print("parameters changed:", changed)

AMP enabled: False
finite loss: True
finite unscaled gradient norm: True
parameters changed: True

On CPU, enabled=False acts as a transparent pass-through: autocast is a no-op, scaler.scale(loss) returns the raw loss, unscale_ does nothing, and scaler.step(optimizer) executes optimizer.step() directly.

Graph compilation after eager verification

torch.compile captures compilable regions through TorchDynamo and sends them to a selected backend. The default backend, TorchInductor, generates optimized kernels and can fuse operations; other backends behave differently.^{[22]Reference 22PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilationhttps://doi.org/10.1145/3620665.3640366[23]Reference 23torch.compilehttps://docs.pytorch.org/docs/2.14/generated/torch.compile.html}

Compilation can improve steady-state performance, but startup cost, graph breaks, shapes, and workload can outweigh the gain. It doesn't repair a faulty loop or a leaked validation split. Check eager correctness first.

The next smoke check uses backend="eager" and fullgraph=True: capture this forward pass as one graph, then execute it with ordinary PyTorch operations. It tests capture and one output comparison, not Inductor performance or an entire training loop. Inductor can also compile CPU code.

import torch
from torch import nn

torch.manual_seed(0)
eager = nn.Linear(2, 2)
compiled = torch.compile(eager, backend="eager", fullgraph=True)
ticket = torch.tensor([[1.2, 0.8]])

eager_logits = eager(ticket)
compiled_logits = compiled(ticket)
print("compiled module:", type(compiled).__name__)
print("logits match eager:", torch.allclose(eager_logits, compiled_logits))

compiled module: OptimizedModule
logits match eager: True

With the default fullgraph=False, graph breaks normally resume Python execution and can cost performance; with fullgraph=True, a break raises an error. For discrepancies, compare seeds and stochastic draws, layer modes, dtypes, tolerances, and backend behavior, then isolate capture versus code generation. Fused arithmetic can differ by rounding. Benchmark steady-state iterations separately from compilation.^{[24]Reference 24Troubleshooting torch.compilehttps://docs.pytorch.org/docs/2.14/user_guide/torch_compiler/torch.compiler_troubleshooting.html}

When combining compiled training with AMP, also check the backward autocast policy. PyTorch 2.14 documents torch._functorch.config.backward_pass_autocast="off" for a compiled region whose backward runs outside autocast. Our eager AMP example and separate forward-only capture check don't exercise that combination.^{[19]Reference 19Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

Artifacts and audit receipts for the training run

A saved weight file doesn't show whether validation was isolated, or which rows produced the loss curve. Keep these records with the run:

Practice: break and repair the loop

Break one step in the runnable examples, and predict the symptom before you run the check:

1. In 03-one-update.py, comment out optimizer.step(). Predict whether changed by step remains True.
2. In 04-zero-grad.py, invoke backward() three times before clearing gradients. Predict the resulting gradient magnitude.
3. In 07b-checkpoint-snapshot.py, replace copy.deepcopy(...) with model.state_dict(). Predict which printed check changes when live weights change.
4. In 08-train-eval-modes.py, leave the model in training mode during the second pair of predictions. Predict whether repeated forward passes match.
5. In 12-amp-ordering-pattern.py, explain why calling scaler.unscale_(optimizer) before clip_grad_norm_ is mandatory.
6. In 13-compile-after-eager.py, modify an eager weight after calling torch.compile and predict whether compiled output changes.

Expected observations

1. Without optimizer.step(), autograd populates .grad, but parameter values remain unchanged (changed by step: False).
2. The gradient accumulates to 6.0: each backward pass adds another derivative of 2.0 into .grad.
3. copied state stayed at zero evaluates to False. Both dictionaries reference live parameter storage, so both follow the weight update.
4. The seeded training-mode pair differs here. Dropout can sample identical masks on repeated passes; inspect model.training to verify the active mode.
5. Clipping scaled gradients applies the wrong unscaled bound, approximately max_norm / S. When a large $S$ makes clipping activate, that can make the gradient much too small; it needn't make it exactly zero.
6. torch.compile wraps the original module rather than copying weights, so in-place edits to eager parameters reflect immediately in compiled outputs.