Switching Between Training, Evaluation, and Inference Modes

Welcome back. Last lesson, you trained an MLP by computing a loss, calling backward(), and letting an optimizer update its parameters. Measuring that model is a different activity: you want predictions and a loss, but you do not want validation data to affect the training procedure.

PyTorch gives you two separate controls for this: a model’s training or evaluation behavior, and autograd’s gradient-tracking behavior. Today you will learn when to switch each one, then put them together in a training-and-validation pattern.

Training and evaluation change how some layers behave

A newly created PyTorch model starts in training mode. Calling model.train() selects that mode explicitly; calling model.eval() selects evaluation mode. The switch matters for layers whose forward-pass behavior depends on the mode. For example, nn.Dropout randomly zeroes some activations during training but leaves them unchanged during evaluation. Batch normalization uses batch statistics and updates its running statistics during training; in evaluation mode, it uses the running statistics learned during training instead.

Watch the selected portion of Lightning AI’s Episode 2: PyTorch Dropout, Batch size and interactive debugging to see why forgetting the switch can change validation predictions.

You can observe the switch without training a whole network:

import torch
from torch import nn

dropout = nn.Dropout(p=0.5)
x = torch.ones(100)

dropout.train()
train_output = dropout(x)

dropout.eval()
eval_output = dropout(x)

print("training output values:", torch.unique(train_output))
print("evaluation leaves input unchanged:", torch.equal(eval_output, x))

During training, the output values will ordinarily include zero and $2$: PyTorch scales the retained values to compensate for the half it randomly drops. During evaluation, eval_output equals x.

Your MLP from the previous lesson contains linear layers and ReLU, but no dropout or batch normalization. Its predictions therefore need not change merely because you call model.eval(). Still, use train() and eval() at the appropriate points: that keeps the code correct if you later add a mode-sensitive layer.

Gradient tracking is a separate decision

The name eval() invites a common mistake: model.eval() does not disable autograd. If the model’s parameters require gradients, a forward pass in evaluation mode can still build a graph for backward(). Conversely, calling model.train() while gradient tracking is disabled does not turn tracking back on.

Read PyTorch’s Autograd mechanics documentation for the distinction. Its discussion of gradient modes also explains why torch.no_grad() and torch.inference_mode() are not quite interchangeable.

In ordinary training code, default gradient mode records the operations needed to differentiate the loss with respect to model parameters. For validation, wrap the forward pass in with torch.no_grad():. Operations inside that block are not recorded for backward, even though the model’s parameters still have requires_grad=True. Leaving the block restores the previous gradient mode.

For computation that will never need to interact with autograd, torch.inference_mode() also disables graph recording and can remove additional tracking overhead. Its restriction is important: tensors created in inference mode cannot subsequently be used in a computation that autograd must record. Use no_grad() as the straightforward default for validation, especially if you may reuse its results in later gradient-tracked work; consider inference mode for self-contained prediction or evaluation code.

These controls answer different questions:

Neither eval() nor a gradient-disabled context erases gradients already stored in parameters’ .grad fields. Neither one updates parameters, either: as you saw last lesson, an update requires an optimizer step.

Put the switches around validation

Continue from the previous lesson’s notebook, where model, X, y, loss_fn, and optimizer are defined. The following three new examples are a small validation set. They illustrate the mechanics; three examples are far too few for a reliable performance estimate.

X_val = torch.tensor([
    [1.5, 0.0, 0.0, 0.0],
    [0.0, 1.5, 0.0, 0.0],
    [0.0, 0.0, 1.5, 0.0],
])
y_val = torch.tensor([0, 1, 2], dtype=torch.long)

model.eval()

# Demonstration only: eval mode by itself still permits graph recording.
tracked_logits = model(X_val)
print("eval output tracks gradients:", tracked_logits.requires_grad)

# The validation pass uses both evaluation behavior and no-grad mode.
with torch.no_grad():
    logits = model(X_val)
    val_loss = loss_fn(logits, y_val)
    val_accuracy = (logits.argmax(dim=1) == y_val).float().mean()

print("validation output tracks gradients:", logits.requires_grad)
print("validation loss:", val_loss.item())
print("validation accuracy:", val_accuracy.item())
print("gradient tracking restored:", torch.is_grad_enabled())

# Before another training update, restore training behavior.
model.train()
optimizer.zero_grad()
train_loss = loss_fn(model(X), y)
train_loss.backward()
optimizer.step()

The first print should be True: evaluation mode did not prevent tracked_logits from being connected to the model’s parameters. The output computed inside no_grad() should report False. After the with block, gradient tracking is restored, so the final training loss can be differentiated as usual.

In a larger program, apply this pattern around each phase: call model.train() before training batches; call model.eval() before validation batches and put their forward passes inside a gradient-disabled context. Do not call backward() or optimizer.step() for validation. Validation examples are for measuring the current model, not fitting it.

If you only need predictions, the validation block can instead use inference mode:

model.eval()
with torch.inference_mode():
    predicted_classes = model(X_val).argmax(dim=1)

print(predicted_classes.tolist())

Treat predicted_classes here as an evaluation result, not as a tensor to feed into a later gradient-tracked computation. If that later use is needed, prefer no_grad().

Conclusion

The reliable distinction is model behavior versus autograd behavior. model.train() and model.eval() select how mode-sensitive layers run; torch.no_grad() and torch.inference_mode() prevent operations from being recorded for backward. For validation, use eval() together with a gradient-disabled context. Before resuming updates, call train() and perform the next forward pass with gradient tracking available.

Next, you will save and restore model and optimizer state so that a training run can continue rather than start over.