Cached features and inference tensors

no_grad and inference_mode both disable reverse-mode recording for ordinary calculations in their blocks. Inference mode can remove additional autograd overhead, but tensors it creates have restrictions outside the block.

The result keeps its inference status

import torch

with torch.inference_mode():
    cached = torch.tensor([1., 2.])
print(cached.is_inference())  # True

w = torch.ones(2, requires_grad=True)
try:
    loss = (cached * w).sum()
except RuntimeError:
    print("Multiplication cannot save the inference tensor")

To compute the gradient with respect to w, multiplication needs cached. It attempts to save that tensor for backward, which inference tensors do not support. The error occurs while building this multiplication, before backward is called.

Not every later operation needs to save the inference tensor. Addition, for example, can work in situations where multiplication fails. The useful rule is to avoid inference outputs where later differentiation may need to save them, rather than assuming every use in a graph fails.

What the mode changes

New inference tensors omit version-counter tracking, and inference mode also removes view-tracking overhead and disables forward-mode AD. This can improve performance when gradients are unnecessary, but the gain depends on the workload. It does not make every operation faster or eliminate all bookkeeping for tensors created before entering the block.

Neither mode switches a module to evaluation behavior. Use model.eval() separately when dropout, BatchNorm or other modules should behave in their evaluation mode.

Ordinary cached features can be constants in a later graph

x = torch.tensor([0., 1.])
with torch.no_grad():
    features = (x * 2).sin()
w = torch.ones(2, requires_grad=True)
(features * w).sum().backward()
print(features.is_inference())  # False
print([round(v, 4) for v in w.grad.tolist()])  # [0.0, 0.9093]

The later loss differentiates with respect to w. It does not recover a gradient path through the earlier feature computation to x. Features are constants in that later calculation.

If you already have an inference tensor, cloning it outside inference mode creates an ordinary tensor with independent data storage. Choosing no-grad when computing features can avoid that extra copy.

def cached_features(x: torch.Tensor) -> torch.Tensor