Skip to content

← Autograd Mechanics step 8 of 8

Medium Primitives

Faster than no_grad, and one-way

with torch.inference_mode():
    preds = model(x)

Everything no_grad does, and two more things:

  • No version counting. The mechanism from the in-place problem is switched off, so every operation skips the bookkeeping.
  • No view tracking. Views made inside do not record the metadata autograd would need to differentiate through them.

Both are pure overhead when nothing will ever be differentiated, which is what makes it the right choice for a serving path.

The catch

Tensors created inside are marked as inference tensors, permanently, and a marked tensor cannot be saved for backward:

with torch.inference_mode():
    cached = expensive(x)

loss = (cached * w).sum()
loss.backward()
RuntimeError: Inference tensors cannot be saved for backward.

The mark does not wear off when the block ends. It is a property of the tensor, not of where it is used.

Choosing between them

  • inference_mode when the results leave via a response and never come back: a serving endpoint, a batch scoring job, a metric you print.
  • no_grad when the results re-enter training code: a target computed from a frozen teacher, an EMA update, a replay buffer, anything a later loss will touch.

When you have an inference tensor and need it in a graph anyway, .clone() outside the block gives an ordinary tensor. That is a real copy, and needing it usually means no_grad was the right choice in the first place.

Your task

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

Compute (x * 2).sin() without building a graph, returning a tensor that a later loss can be differentiated through. x itself does not require gradients; the result must simply not be an inference tensor.