Lesson 10 / 33

no_grad, inference_mode, detach and zeroing gradients

Turn gradient tracking off where it is not needed, detach tensors safely, and reset gradients every training step.

By · 15 years in web development

When not to track gradients

Recording operations costs memory, because autograd keeps intermediate values for the backward pass. Whenever you will not call backward(), switch tracking off.

with torch.no_grad(): runs the enclosed code without building a graph. Use it for validation, for inference, and for manual weight updates (an update must not itself become part of the graph). torch.inference_mode() is a stricter and slightly faster version: tensors created inside it can never be used in autograd later, which is fine for pure inference.

tensor.detach() returns a tensor that shares the same data but is cut off from the graph. Use it when you need a value for logging or plotting, or as a fixed target that gradients should not flow into.

Finally, because .grad accumulates, each training step must clear old gradients first. optimizer.zero_grad() does this; by default it sets gradients to None, which saves memory. Accumulating on purpose over several small batches before one step() is a known trick to simulate a bigger batch.

A manual update step

This is what an optimizer does internally. The update runs under no_grad, and the gradient is reset every iteration. Working it by hand: w goes from 1 to 1.9, then 1.99, then 1.999, approaching the answer 2.

import torch

w = torch.tensor(1.0, requires_grad=True)
for step in range(3):
    loss = (w * 3.0 - 6.0) ** 2       # minimum at w = 2
    loss.backward()
    with torch.no_grad():
        w -= 0.05 * w.grad            # update must not be recorded
    w.grad = None                     # reset, or gradients accumulate
    print(step, round(w.item(), 4))

Off the record

torch.no_grad() is like telling a journalist "this part is off the record". You still say everything and the conversation still happens, but nothing is written into the notebook, so nobody can trace it back later. That is exactly what you want during an exam-style evaluation of your model.

Two classic memory and correctness bugs

Forgetting zero_grad() makes each update use the sum of all past gradients, so training diverges or behaves strangely. Running validation without no_grad() or inference_mode() builds graphs you never use, and on a GPU that often ends in an out-of-memory error halfway through evaluation.

Quick check: You want to record the loss value every step for a plot, without keeping the computation graph alive. What should you store?

  • The `loss` tensor itself, appended to a list
  • `torch.tensor(loss, requires_grad=True)`
  • `loss.item()`
  • The result of calling `loss.backward()`
Answer

`loss.item()` — `loss.item()` returns a plain Python float with no link to the graph. Storing the tensor keeps the whole graph and its memory alive.