🎓 A1Learner
🔍

PyTorch Inference Mode

Inference mode is a context manager that switches autograd off completely for a block of code. Tensors created inside it are inference tensors, and strict rules guard them.

Evaluation under no_grad

no_grad stops gradient tracking for the block. The output stays a normal tensor, and it can slip into tracked maths later with no warning.

import torch

model = torch.nn.Linear(4, 2)
x = torch.randn(2, 4)
with torch.no_grad():
    out = model(x)
print(out.shape, out.requires_grad)

The same evaluation in inference mode

Inference mode runs the same forward pass. A tensor created inside the block is marked as an inference tensor.

with torch.inference_mode():
    out = model(x)
    kept = x * 2
print(out.shape)
print(kept.is_inference())

The rails bite outside the block

Using the inference tensor in tracked maths raises an error at once. The mistake surfaces at the boundary, not deep inside a training step.

w = torch.randn(4, requires_grad=True)
try:
    kept @ w
except RuntimeError:
    print("RuntimeError raised")

Clone makes a normal tensor

A clone of an inference tensor is a normal tensor. The clone can take part in tracked maths.

normal = kept.clone()
print(normal.is_inference())
print((normal @ w).shape)

Example

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Linear(4, 8),
    torch.nn.ReLU(),
    torch.nn.Linear(8, 3),
)
model.eval()
x = torch.randn(5, 4)

with torch.no_grad():
    out_no_grad = model(x)

with torch.inference_mode():
    out_inference = model(x)
    saved = x * 2

print("Output shape:", out_inference.shape)
print("Outputs equal:", torch.equal(out_no_grad, out_inference))
print("Saved tensor is inference:", saved.is_inference())

w = torch.randn(4, requires_grad=True)
try:
    y = saved @ w
    print("No error raised")
except RuntimeError as e:
    print("Error type: RuntimeError")
    print("Error message starts:", str(e).split(".")[0])

Key points

  • Inference mode: a block where autograd records nothing at all
  • Inference tensor: a tensor created inside inference mode
  • Use torch.inference_mode() around evaluation and prediction code
  • Inference tensors cannot be saved for backward in a later computation
  • In-place edits to an inference tensor outside the block raise an error too
  • Use clone() on an inference tensor to get a normal tensor
  • In short, inference mode swaps flexibility for a leaner and safer evaluation

In this example, a small Sequential model scores five rows of four values. The first pass runs under no_grad and the second under inference mode. Both passes print the same output shape, and the outputs compare equal. The tensor saved inside inference mode reports True for is_inference(). Outside the block, that tensor meets a weight that requires gradients. PyTorch raises a RuntimeError at that line. The program prints only the first sentence of the error, because later sentences add advice that can change between versions. A clone of the saved tensor would have worked, at the cost of one copy.

Note

torch.no_grad() is the softer switch. Its tensors stay normal and can rejoin tracked work later.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What happens when an inference tensor is used in a computation that autograd tracks?