The Training Loop in Python

Scikit-Learn hides the entire optimization process behind one method call: model.fit(X, y). PyTorch deliberately does not offer an equivalent — training a nn.Module means writing the loop yourself, batch by batch, epoch by epoch.

That's not an oversight; it's the trade-off PyTorch makes for research flexibility. Custom loss functions, unusual architectures, multiple optimizers, gradient clipping, mixed precision — all of these are just extra lines inside a loop you already control, rather than configuration flags bolted onto a black-box .fit().

The loop itself always follows the same five-step shape regardless of what model you're training: a forward pass to get predictions, a loss calculation to measure error, then the three-step optimization sequence of zeroing gradients, computing new ones with backward(), and applying them with an optimizer step. The next sections walk through each step in order.

Iterating a DataLoader — for X_batch, y_batch in dataloader: — hands you one mini-batch of inputs and labels at a time rather than the whole dataset at once, which keeps memory usage bounded regardless of how large the underlying dataset is.

The forward pass is the line predictions = model(X_batch). Calling a module like a function invokes its forward() method under the hood, running X_batch through every layer in sequence — matrix multiplications, activations, whatever the architecture defines — to produce raw output predictions (often called logits for classification).

Nothing about the model's weights changes during this step. The forward pass only computes an output; PyTorch's autograd engine is quietly building a computation graph behind the scenes as it goes, which is what makes the later backward pass possible, but no learning has happened yet.

To answer the exercise directly: during the forward pass, the input batch's data literally flows through the network's layers — each nn.Linear, convolution, or activation function transforms the tensor in sequence — and what comes out the other end is the model's current prediction for that batch, given its current weights.

It's easy to conflate this with 'training' happening, but the forward pass by itself is just inference: run the same input through the same weights twice in a row (without any update in between) and you get the identical output both times. The result only becomes useful for learning once it's compared against the true labels in the next step.

That comparison — and the number it produces — is exactly what Step 2, the loss calculation, covers next.

Step 2 is loss = loss_fn(predictions, y_batch) — comparing what the model predicted against the true labels for that batch, using a loss function like nn.CrossEntropyLoss for classification or nn.MSELoss for regression.

The result, loss, is a single scalar tensor, but it's not just a plain number — it also carries a reference to the entire computation graph that produced it, all the way back through the forward pass to the model's parameters. That graph is what autograd needs in the next step to figure out how each individual weight contributed to the error.

Different loss functions encode different notions of 'wrong': CrossEntropyLoss penalizes confident incorrect class predictions heavily, while MSELoss penalizes large numeric distances between a predicted and actual value. Picking the right loss function for the problem is as important as picking the right architecture.

The purpose of the loss function in Step 2 is to turn 'how wrong was the model' into a single differentiable number. Without that conversion, there'd be nothing for autograd to compute a gradient of — accuracy alone, for instance, isn't smooth or differentiable, which is exactly why loss functions like cross-entropy exist as differentiable proxies for the metric you actually care about.

That scalar loss value is also what most training scripts log and plot per batch or per epoch — a steadily decreasing loss curve is the most direct signal that the optimization is working, while a flat or exploding curve usually points to a bug (wrong learning rate, forgotten zero_grad(), or a mismatched loss function for the task).

Critically, the loss function does not touch the model's weights itself. It only produces the error signal; the actual parameter updates happen in Steps 3 through 5, covered next.

Steps 3 through 5 are the trio that actually updates the model: optimizer.zero_grad(), loss.backward(), optimizer.step() — and the order matters. zero_grad() clears out the .grad attribute on every parameter before this batch's gradients are computed, because PyTorch accumulates gradients into .grad by default rather than overwriting them; skipping this step silently mixes gradients from the previous batch into the current one.

loss.backward() walks the computation graph the loss tensor carries with it and, via the chain rule, computes d(loss)/d(weight) for every parameter that requires gradients, storing each result in that parameter's .grad. Nothing about the weights changes yet — this step only computes 'which direction reduces the loss.'

optimizer.step() is what actually moves the weights, using whatever update rule the optimizer implements (plain SGD, Adam, etc.) applied to the gradients .backward() just populated, scaled by the learning rate.

optimizer.step() doesn't recompute anything about the loss or the forward pass — it only reads the .grad tensors that loss.backward() already populated on every registered parameter, and applies its update rule to nudge each weight a small step in the direction that reduces the loss, scaled by the learning rate.

Which exact update rule runs depends on which optimizer you instantiated: plain SGD moves each parameter directly opposite its gradient, while Adam maintains running estimates of the gradient's mean and variance to adapt the effective step size per parameter. The training loop code itself — zero_grad(), backward(), step() — stays identical either way; only the optimizer object changes.

Because step() depends entirely on .grad being fresh and correct, calling it without a preceding backward() in that iteration (or after gradients were left dirty from a skipped zero_grad()) is what produces the classic symptom of a broken training loop: a loss that barely moves or that oscillates erratically.

The five-step loop covers how a single batch gets trained on, but there's a mode-switching detail that sits outside that loop entirely and trips up almost everyone the first time: some layers behave differently depending on whether the model thinks it's currently training or being evaluated.

That distinction doesn't show up anywhere in zero_grad, backward, or step — it's a separate flag on the model itself, set once before you start iterating, and it has to be set correctly or the model's outputs become unreliable in ways that have nothing to do with the optimizer.

The next two sections cover exactly which layers care about this and the one-line call that controls it.

Dropout and BatchNorm are the two layers this matters for most. Dropout randomly zeroes out a fraction of activations during training as a regularization technique — but you obviously don't want random neurons disabled when the model is actually making a real prediction, so during evaluation it should act as a no-op and pass every activation through unchanged.

BatchNorm behaves differently in a subtler way: during training it normalizes each batch using that batch's own mean and variance, while during evaluation it switches to running statistics accumulated across all of training, since a single inference batch (sometimes even a batch of one) doesn't have a reliable mean and variance of its own.

PyTorch has no way to infer which behavior you want automatically — the model doesn't know whether the current call is part of a training step or a real prediction unless you tell it explicitly, which is exactly what the next section's method does.

model.train() is the call that flips the switch: it sets self.training = True on the module and recursively on every submodule, and layers like Dropout and BatchNorm check that exact flag internally to decide which behavior to run. It's a lightweight, instant call — no computation happens, it's purely a mode flag.

Its counterpart, model.eval(), sets that same flag to False everywhere, which is what you call before running validation during training or before serving real predictions after loading a saved model. Forgetting to switch back to model.train() after a validation pass in the middle of a training loop is a common, quiet bug — the next epoch would keep training with Dropout and BatchNorm stuck in evaluation behavior.

A new nn.Module defaults to training mode, so model.train() isn't strictly required the very first time, but calling it explicitly at the top of every training loop is the safe habit — especially in any script that also runs periodic validation with model.eval() in between.

Put the whole loop together and a single training epoch looks like: model.train(), then for each batch — forward pass, compute loss, zero_grad(), backward(), step() — and if a validation pass runs afterward, switch to model.eval() (typically wrapped in torch.no_grad() to skip building an autograd graph you won't use) before switching back to model.train() for the next epoch.

Every one of these five steps and the train/eval switch matters independently: skip zero_grad() and gradients silently accumulate across batches; skip model.eval() during validation and Dropout/BatchNorm corrupt your metrics; skip backward() or step() and the model simply never learns despite the loop running without errors.

With a trained model in hand, the natural next question is how to persist it past the current Python process — which is exactly what model saving and loading, the next lesson, covers.

In Scikit-Learn, training is just model.fit(X, y). In PyTorch, you have to write the training loop from scratch. This gives you absolute control.

Step 1: The Forward Pass. You feed the batch of data (X) into the model to get the predictions.

What happens during the "Forward Pass" of the training loop?

• →The network updates its weights.
• →The input data flows through the network's layers to produce an output prediction.
• →The error is calculated.

Step 2: Calculate the Loss. You compare the model's predictions against the true answers (y) using a Loss Function (like CrossEntropyLoss).

What is the purpose of the "Loss Function" in Step 2?

• →To increase the learning rate.
• →To mathematically quantify how wrong the model's predictions were compared to the actual target labels.
• →To generate the confusion matrix.

Steps 3, 4, and 5 form the core of Optimization. Zero the gradients, calculate the new gradients (backward), and update the weights (step).

In the final optimization sequence (zero_grad, backward, step), what exactly does optimizer.step() do?

• →It moves the data to the next CPU core.
• →It takes the gradients calculated by .backward() and uses them to slightly adjust the model's weights to reduce the error for the next loop.
• →It calculates the accuracy score.

Now, prepare yourself. We are about to enter the ADA Defense Protocol. Ensure you understand training modes vs evaluation modes.

Certain neural network layers (like Dropout and BatchNorm) behave differently during Training vs Testing. You must explicitly tell the model its current state.

ADA DEFENSE: Before starting your for loop to train the network, what PyTorch method MUST you call on the model to activate layers like Dropout?

• →model.fit()
• →model.train()
• →model.eval()

Threat neutralized. Model states verified. Proceeding to Model Saving and Deployment.

Run One Real Training Step. Finish training_step(): apply the gradient descent update to the weight.