Enable training-specific layer behavior
Training
PyTorch · Module
model.train
Sets a module and its descendants to training mode. This affects layers with training-specific behavior, but it does not control gradient tracking.
How it works
Calling train() switches this module and its child modules into training mode by setting their training flag. Modules such as dropout and batch normalization can use that flag to change behavior; the method does not itself enable or disable gradient computation.
Syntax
Module.train(mode=True)Parameters
| Name | Type | Description |
|---|---|---|
mode | bool = True | Whether to set the module to training mode; false selects evaluation mode. |
Returns
Returns the module itself.
Key points
- A newly created module starts in training mode; train() explicitly selects that state.
- Use train() at the start of a training phase.
- Use eval(), equivalent to train(False), for validation, testing, or inference rather than train().
- The recommended phase transition is training mode for optimization and evaluation mode for measuring performance.
Common mistakes
- Do not use train(False) or eval() as a replacement for disabling autograd; use an appropriate gradient context when gradients must be disabled.
- Training/evaluation mode and gradient/inference mode are independent controls.
- Explicitly select the correct mode when using layers whose behavior depends on training state.
Example: Enable training mode
import torch
from torch import nn
model = nn.Sequential(nn.Linear(4, 8), nn.Dropout())
model.train()
output = model(torch.randn(2, 4))Creates a model, enables training behavior, and performs a forward pass.
When to use it
Switch modes explicitly between phases
Validation or testing
Interview questions
QWhat does Module.train() do?
train() puts the module into training mode and propagates that state to its child modules. This matters for components such as dropout and batch normalization, which may behave differently during training.
Follow-up: Does train() turn on gradients? No. Mode selection is independent of gradient tracking, so use no_grad() or inference_mode() when appropriate.
QWhat is the relationship between train(False) and eval()?
The optional boolean controls the mode: train(True) selects training mode, while train(False) is equivalent to eval(). The chosen state is applied recursively to the module hierarchy.
Follow-up: Why can’t eval() replace no_grad()? Because evaluation mode does not disable gradient computation; those mechanisms control different behaviors.
QWhen should train() be called?
Call model.train() for the training loop and model.eval() for validation or testing. This ensures state-sensitive layers use the behavior intended for the current phase.
Follow-up: What failure can occur if eval() is omitted during validation? Dropout can remain active and batch-normalization running statistics can be updated using validation data.
Related
torch.nn.Module.evaltorch.no_gradtorch.inference_modetorch.nn.Dropouttorch.nn.BatchNorm2d