Implementing a Typical Training Loop for Neural Networks

A training loop repeats the same update round many times.

Programming

Testing the Belief About model.eval()

Many learners who first meet model.eval() come away believing it is a single switch that turns learning off and makes the model ready for predictions. The belief sounds reasonable and is partly close to the truth, which is exactly why it causes confusion. This article starts from that belief, tests it, and then rebuilds an accurate picture of what the two mode calls do and what they do not do.

What do you think happens?

Does calling model.eval() by itself also stop the model from tracking gradients?

Yes, model.eval() turns off gradient tracking
No, gradient tracking is handled by a separate context such as torch.no_grad()

Reveal answer

Answer:

No, gradient tracking is handled by a separate context such as torch.no_grad()

The source treats the mode setting and the no-gradient context as two different pieces. model.eval() is about how components such as dropout and batch normalization behave, while torch.no_grad() is the piece that disables gradient tracking. The source uses them together rather than interchangeably.

Tracing the Order of One Update Round

A training loop repeats the same update round many times. Each round computes a loss from the model's output, works out how the parameters should change, and then applies that change. The source places the gradient reset at the start of the round, so the steps of one round run in this order: reset the gradients, run the forward pass, compute the loss, call backward to compute gradients, and call the optimizer step to update the parameters.

What is the sequence of steps in one iteration, and where does the gradient reset sit relative to backward() and step()?

Read the diagram as a sequence of hand-offs. The forward pass and loss come first because they produce the quantity that backward() differentiates. The backward call fills in gradients in the computation graph that PyTorch built in the background while the forward pass ran. The optimizer step then uses those gradients to change the parameters so that the loss goes down. For the SGD optimizer, the source describes this as multiplying the gradients by the learning rate and adding the scaled negative gradient to the parameters.

What Train and Eval Modes Change

The source introduces model.train() and model.eval() as settings that place the model into a training mode or an evaluation mode. The reason given is that some components of a network behave differently during training than during inference. Dropout and batch normalization are the examples named. The mode setting is therefore about how those components behave. It is not a general on-off switch for everything the model does, and it does not by itself decide whether gradients are computed.

What state changes when a model moves between training and evaluation mode?

The diagram shows the mode as a state the model can be in, with one call moving it to the other state. It says nothing about gradients. Gradient handling is treated by the source as a separate concern that appears elsewhere in the loop. Keep that separation in mind, because it explains why a prediction step needs more than one line.

python

for epoch in range(num_epochs):
    model.train()
    for inputs, targets in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
        loss.backward()
        optimizer.step()

for epoch in range(num_epochs):
    model.train()
    for inputs, targets in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
        loss.backward()
        optimizer.step()

Notes on the code

model.train() is placed before the update rounds so the model is already in training mode when they begin
optimizer.zero_grad() opens each round
loss.backward() and optimizer.step() form the second half of the round

This skeleton is an illustrative arrangement of the calls the source describes, not a listing taken from it. Its purpose is to show where the mode call sits relative to the update rounds. If the architecture later gains a dropout layer, the model is already set to training mode at the start of each epoch, so the new layer behaves as intended.

Where zero_grad Sits and What Accumulation Does

The source states that an optimizer.zero_grad() call belongs in each update round so that gradients are set back to 0 before new ones are computed. The cost of leaving it out is accumulation. Without the reset, the gradients produced by one round are still present when the next round runs, and the new gradients are added to them. The source notes that this accumulation may be undesired, which is the reason the reset is part of the standard loop.

What happens to the stored gradient values across successive update rounds if zero_grad() is skipped?

The diagram illustrates the accumulation described in the source. If each round's gradients are never cleared, the stored values carry forward and grow, so the update applied in a later round reflects more than the current batch. Calling optimizer.zero_grad() at the start of every round prevents this, because each round then begins from zeroed gradients and computes only its own.

Before

for inputs, targets in train_loader:
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()

No reset: gradients from earlier rounds remain and are added to new ones

After

for inputs, targets in train_loader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()

Reset at the start of each round: gradients begin at 0 before new ones are computed

Making Predictions After Training

The source introduces the prediction step with the sentence that after training, the model can be used to make predictions. That is the only role it plays. The step follows training and does not perform learning. The source's version enters evaluation mode, disables gradient tracking with torch.no_grad(), runs the model on the training inputs, and prints the result. Each line has its own effect, so the trace below follows them one at a time.

python

model.eval()
with torch.no_grad():
    predictions = model(train_inputs)
print(predictions)

model.eval()
with torch.no_grad():
    predictions = model(train_inputs)
print(predictions)

Notes on the code

model.eval() changes the mode of components such as dropout and batch normalization
torch.no_grad() is a separate context that stops gradient tracking for the enclosed computation
The model runs inside the context, so no gradient graph is built for these outputs
print(predictions) displays the result of the prediction run

Worked example

Tracing one prediction run

Trace what each line of the prediction step contributes, given that the model has just finished training.

1
Set evaluation mode
model.eval() moves the model into evaluation mode, so layers whose behavior differs between training and inference, such as dropout and batch normalization, use their evaluation behavior.
2
Open the no-gradient context
torch.no_grad() tells PyTorch not to track gradients for the computation inside the with block. This is a separate piece from the mode call.
3
Run the model
The model processes the inputs inside the context. Because gradient tracking is off, the outputs are computed without a gradient graph being built for them.
4
Print the result
The predictions are displayed. Nothing in this step calls backward() or optimizer.step(), so the parameters do not change.

Result

The model produces predictions in evaluation mode with gradient tracking disabled. Each line contributes a different effect, and removing either the mode call or the no-gradient context changes what the step does.

Mistakes Beginners Make with the Loop

Common mistake
Treating model.eval() as a complete switch that turns off learning and gradient tracking
What it looks like
Calling model.eval() once and then running predictions without any context manager, assuming gradients are no longer computed.
Why it fails
The source describes model.eval() as placing the model into an evaluation mode for components such as dropout and batch normalization. Gradient tracking is treated as a separate concern.
Do this instead
Use model.eval() for the mode, and wrap the prediction in torch.no_grad() to disable gradient tracking.
Common mistake
Omitting optimizer.zero_grad() from the update round
What it looks like
Placing backward() and step() in the loop without a reset, so each round starts with gradients left over from earlier rounds.
Why it fails
The source states that without the reset, gradients accumulate across rounds in ways that may be undesired.
Do this instead
Call optimizer.zero_grad() at the start of each update round, before the new gradients are computed.
Common mistake
Assuming the mode calls are required for a network without dropout or batch normalization
What it looks like
Believing that a network with only plain layers is affected by model.train() or model.eval().
Why it fails
The source notes that its NeuralNetwork class has no dropout layer or other mode-affected component, so the calls do nothing for that model.
Do this instead
Still include the mode calls. They are harmless for such a model and keep the code correct if the architecture changes later.
Common mistake
Placing the mode call inside the inner loop without a clear reason, which obscures where the mode is set for the whole training run.

Practice Tracing the Loop

MEDIUM

Suppose a colleague moves zero_grad() to the end of each round, after optimizer.step(). Explain whether the training loop still resets gradients before new ones are computed, and identify the step in the round that the reset must come before.

Hints

Look at the order the source describes: the gradients are set back to 0 before new ones are computed.
Consider which call produces the new gradients, and which call applies them to the parameters.
A reset placed after step() still clears the gradients, but think about when the next backward() call happens relative to it.

Your attempt (not saved)

EASY

A model has a dropout layer and is used to make predictions on a validation set. Describe which mode it should be in for that prediction run, and which context manager should wrap the prediction, and explain what each one is responsible for.