PyTorch
model training
eval mode
deep learning
machine learning

How to check if a model is in train or eval mode in PyTorch?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

In PyTorch, a module is either in training mode or evaluation mode, and the difference matters for layers such as Dropout and BatchNorm. The quickest way to check the current state is to inspect model.training, which is a boolean flag updated when you call model.train() or model.eval().

The Simple Check

Here is the direct pattern:

python
1import torch.nn as nn
2
3model = nn.Sequential(
4    nn.Linear(10, 20),
5    nn.Dropout(p=0.5),
6    nn.BatchNorm1d(20),
7)
8
9print(model.training)  # True by default
10
11model.eval()
12print(model.training)  # False
13
14model.train()
15print(model.training)  # True

A newly created module starts in training mode, so model.training is normally True until you switch it.

What The Mode Actually Changes

The flag does not affect every layer equally. It mainly changes behavior for modules that need distinct training and inference semantics.

Examples:

  • Dropout randomly drops activations during training and becomes a pass-through during evaluation.
  • BatchNorm uses batch statistics during training and running statistics during evaluation.

That is why forgetting model.eval() before validation or inference can produce noisy or unstable predictions.

The Mode Propagates To Submodules

Calling model.train() or model.eval() applies recursively to child modules as well. In other words, you usually do not need to switch each layer manually.

You can inspect submodules too:

python
for name, module in model.named_modules():
    print(name, module.training)

This is helpful when debugging custom models that mix built-in layers with your own submodules.

eval() Is Not The Same As no_grad()

A common source of confusion is that model mode and gradient recording are separate concepts.

model.eval() changes module behavior.

torch.no_grad() disables gradient tracking for the operations executed inside its context.

You often want both during inference:

python
1import torch
2
3model.eval()
4
5with torch.no_grad():
6    output = model(input_tensor)

Using only one of them solves only part of the inference setup.

Custom Modules Can Use self.training

If you write your own nn.Module, you can branch behavior based on the same flag:

python
1import torch
2import torch.nn as nn
3
4class MyLayer(nn.Module):
5    def forward(self, x):
6        if self.training:
7            return x + 1
8        return x

That keeps custom layers aligned with the rest of PyTorch’s train-versus-eval convention.

This becomes especially important in validation loops. A common training pattern is to switch to evaluation mode for validation, run inference under torch.no_grad(), and then switch back to training mode before the next optimization step. Making that mode transition explicit prevents subtle state bugs from accumulating across epochs.

If you build custom modules, exposing the mode through self.training keeps your code aligned with the rest of PyTorch. That way, one top-level model.eval() call can switch both built-in layers and your own custom logic into inference behavior consistently.

Inspecting a specific submodule can also be useful during debugging. If a complex model wraps Dropout or BatchNorm deep inside nested blocks, printing module.training on those submodules confirms whether the mode change propagated the way you expect through the hierarchy.

That quick check is often enough to catch a forgotten model.train() or model.eval() call during experimentation.

Common Pitfalls

One common mistake is checking the mode on the wrong object. The relevant flag lives on the module instance, so inspect model.training or the specific submodule you care about.

Another mistake is assuming torch.no_grad() implies evaluation mode. It does not change Dropout or BatchNorm behavior.

A third issue is calling model.eval() for validation and then forgetting to switch back to model.train() before continuing training.

Summary

  • Check model.training to see whether a PyTorch module is in training or evaluation mode.
  • Use model.train() to enable training behavior and model.eval() to enable inference behavior.
  • The mode propagates through submodules automatically.
  • 'model.eval() and torch.no_grad() solve different problems, and inference often needs both.'

Course illustration
Course illustration

All Rights Reserved.