PyTorch Training Loops

The previous chapter asked whether a model-system change helps customers. This chapter answers an earlier question: how does a model acquire new behavior at all?

Suppose an e-commerce support queue receives refund tickets. Each ticket has two simple signals: urgency score and negative-sentiment score. A label says whether a specialist later judged that ticket to need escalation. At first, a small classifier guesses badly. A training loop repeatedly shows it examples, measures wrongness, computes which parameters contributed to that wrongness, and updates those parameters.

PyTorch makes that process explicit. That matters when you later fine-tune a text classifier or read a transformer training script: wrappers may add logging and distributed execution, but learning still comes from batches, logits, loss, gradients, optimizer steps, validation, and saved model state.^[1][2]

By the end, you will be able to:

• Build and train a small nn.Module with TensorDataset and DataLoader.
• Explain why classification loss receives raw logits and integer labels.
• Distinguish gradient computation from parameter updates.
• Measure held-out performance without accidentally training on validation rows.
• Save and reload a checkpoint, then diagnose common loop failures.

Turn a ticket into tensors

A model can't consume the phrase "angry customer with delayed refund" directly in this small lab. Pretend an earlier feature pipeline already produced two normalized numbers:

Each training row becomes a vector of two floating-point inputs and one integer class target. A batch stacks those rows into an input tensor. Six tickets have feature shape (6, 2) and label shape (6,).

float32 inputs can flow through linear layers. torch.long labels are class indices for this classification objective. Keeping the shape and dtype contract visible prevents a large class of training bugs.

Logits are scores, not probabilities

A classifier ending in nn.Linear(2, 2) produces two scores for each ticket:

These scores are logits. They may be negative and they don't have to sum to one. For ordinary multi-class classification, PyTorch's CrossEntropyLoss is designed to receive those raw logits plus target class indices. It handles the log-softmax calculation in a numerically stable form.^[2]

Failure case: Don't call softmax() and then feed those probabilities into CrossEntropyLoss for this setup. Pass raw logits. A hand-written log(softmax(logits)) calculation is also easier to make numerically unstable when a score becomes very large.

One update has five actions

For one mini-batch, training follows a fixed order:

Backpropagation is credit assignment: if raising one weight would have increased loss, its gradient points in the direction the optimizer should avoid. Autograd performs this chain-rule bookkeeping for every parameter used in the forward computation.^[3][1]

The following lab freezes a simple linear model into known starting weights, runs exactly one update, and proves that backward() fills gradients before step() moves weights.

If you can point to the line where model state changes, training scripts stop looking like rituals.

Why clearing gradients isn't optional

PyTorch adds newly computed gradients into the existing .grad fields. This enables deliberate gradient accumulation, where several micro-batches contribute to one larger effective update. In a normal one-update-per-batch loop, stale gradients are accidental.

Here the same derivative is silently doubled when it isn't cleared. optimizer.zero_grad() is the standard loop form; setting parameter gradients to None expresses the same reset idea in this minimal demonstration.

Mini-batches make a loop

A dataset owns examples. A data loader groups them into mini-batches and optionally changes row order between epochs. A mini-batch is large enough to estimate a useful gradient but small enough to fit memory.

Notice the final mini-batch contains one row. When you average epoch loss, weight each batch loss by its row count so a short last batch doesn't count as much as a full one.

Build the smallest training loop

Our small model has a hidden layer so it resembles the structure of larger neural classifiers, while staying quick enough to inspect. Before measuring held-out behavior, start with an essential diagnostic: can it memorize one clean batch?

An overfit-one-batch test isn't a product metric. It is a wiring test. If a model can't drive training loss down on eight separable rows, investigate labels, shapes, loss, optimizer order, and learning rate before you scale up.

Loss should fall sharply and memorized batch should be true. Only now is it sensible to ask whether training generalizes.

Separate training from validation

Training rows are allowed to change weights. Validation rows are held out from updates and answer a different question: does the current model work on unseen labeled tickets?

Use two switches during validation:

The next complete loop trains on eight tickets and selects a checkpoint using four held-out tickets. The data is deliberately easy so you can focus on mechanics rather than feature engineering.^[2]

copy.deepcopy matters here. A plain best_state = model.state_dict() would keep references to tensors that later epochs continue to change. Saving immediately with torch.save() is another reliable checkpoint-selection pattern.

Modes change model behavior

The linear classifier would produce the same result in train() and eval() because it contains no mode-sensitive layer. Real models often contain Dropout. During training, Dropout randomly hides activations to regularize learning. During evaluation, it must be disabled so one ticket receives stable scores.

This is why calling only torch.no_grad() isn't enough for evaluation: gradient tracking and layer behavior are separate concerns.

Save evidence you can reload

Training isn't complete when the process exits. You need a portable artifact and enough configuration to reproduce how it was created. For inference, saving a model's state_dict is the standard flexible pattern; for resuming training, also save optimizer state and current epoch.

In production, save the feature schema, label mapping, split or dataset version, model architecture configuration, and metric used to choose the checkpoint. A file of weights without that receipt is difficult to audit.

Debug failures before scaling

When a loop runs without crashing, it can still be wrong. Debug in a narrow order:

Gradient clipping doesn't fix bad data or a broken loss, but it can cap an extreme update once you have identified that gradient norms are the issue.

Numerical triage when loss becomes NaN

When loss becomes nan or inf, locate the first invalid number instead of immediately lowering the learning rate. Test inputs first, then logits, then loss, then gradients after backward(), then parameters after step(). That sequence tells you whether failure entered through data, unstable arithmetic, gradient scaling, or an oversized update.

clip_grad_norm_() returns the total norm before clipping. In a live loop, pass error_if_nonfinite=True as above so a non-finite total norm raises before it can enter optimizer state. Also check torch.isfinite(xb), torch.isfinite(logits), torch.isfinite(loss), and parameter gradients at the first failing step. Find where bad numbers appear before changing hyperparameters blindly.

Inference shouldn't build training history

When serving the selected ticket router, you aren't going to call backward(). torch.inference_mode() disables gradient tracking and additional autograd bookkeeping for computations that won't re-enter a gradient-tracked training graph. Use torch.no_grad() during simple validation when you want the familiar training-script pattern; use inference_mode() for committed inference paths after testing that contract.

Optional GPU extension: automatic mixed precision

The loop above is the core contract. On compatible accelerators, automatic mixed precision (AMP) may reduce memory use and increase throughput by letting suitable operations use lower precision while preserving higher precision where needed. With float16 training, gradient scaling through GradScaler reduces the risk that small gradients underflow before the update. PyTorch's AMP examples place autocast around forward and loss computation, scale before backward(), unscale before clipping, then step and update the scaler.^[4][5]

GradScaler.step() checks for non-finite gradients after unscaling and skips the optimizer update when it finds them. GradScaler.update() then adjusts the scale for later iterations. Treat AMP as an extension after a full-precision loop passes its correctness checks. Faster incorrect training is still incorrect training.

A training receipt for the next chapter

This chapter made a model learn from already prepared tensors. A real ticket-routing model still needs evidence for where its labeled rows came from, whether duplicates leaked across splits, and which preprocessing produced those two features.

Before declaring a training run acceptable, retain:

Practice: Break and repair the loop

Use the runnable examples above as a controlled failure lab. Make one change at a time, predict the result, then run the example and explain the observed behavior.

1. In 03-one-update.py, comment out optimizer.step(). Predict whether changed by step remains True.
2. In 04-zero-grad.py, call backward() a third time before clearing gradients. Predict the gradient value.
3. In 07-train-and-validate.py, replace copy.deepcopy(model.state_dict()) with model.state_dict(). Explain why best_state no longer preserves the best epoch after later updates.
4. In 08-train-eval-modes.py, leave the model in training mode for the second pair of predictions. Predict whether the outputs must match.
5. In 12-amp-ordering-pattern.py, explain why scaler.unscale_(optimizer) belongs before gradient clipping.

Expected observations

1. Without optimizer.step(), gradients are populated but parameter values don't change.
2. The gradient becomes 6.0: each backward() adds another derivative of 2.0 until gradients are cleared.
3. A plain state_dict() result keeps references to parameter storage. Later training updates overwrite the intended snapshot.
4. Dropout remains active in training mode, so repeated outputs can vary. A matching pair can still occur by chance.
5. AMP scales gradients before backward(). Clipping before unscale_() would apply the threshold to scaled values instead of the gradients used for the real update.

Explain the loop without looking back

Explain the complete route from one ticket batch to a saved model without using the word "magic." Then answer these checks.

Evaluation rubric

• Foundational: Identifies logits, loss, gradients, and optimizer updates, including the difference between backward() and step().
• Intermediate: Builds a runnable mini-batch loop with held-out validation, correct modes, and a reloadable checkpoint.
• Advanced: Diagnoses gradient, numerical, or split-quality failures and explains when AMP or clipping belongs in a validated training system.

Common Pitfalls

• Applying softmax() before CrossEntropyLoss instead of giving the loss raw logits.
• Calling backward() batch after batch without clearing gradients or intentionally scheduling accumulation.
• Reporting validation while the model remains in training mode or while held-out tickets influence updates.
• Serving a checkpoint without its label mapping, feature schema, split identity, and selection metric.

Reuse this training pattern

• A small ticket-escalation classifier with an explicit PyTorch optimization loop.
• A one-batch diagnostic that fails early when a new loop is wired incorrectly.
• A validation-and-checkpoint path that selects weights without learning from held-out tickets.
• A training receipt that the data-quality pipeline can make reproducible.

Key Terms

• Logit: Raw class score emitted before a probability normalization.
• Loss: Scalar objective measuring prediction error for optimization.
• Autograd: PyTorch system that computes derivatives through recorded operations.
• Gradient: Directional derivative stored for a parameter after backward().
• Mini-batch: Group of examples used for one gradient estimate.
• Epoch: One traversal through training examples.
• Checkpoint: Saved model state, often paired with configuration and training metadata.