The value backward needed

Some backward computations need values saved during the forward pass. Overwriting one of those values can make backward fail, even when the forward calculation returned the numbers you expected.

import torch

x = torch.tensor([1., 2.], requires_grad=True)
y = x.exp()
y[0] = 0.0
try:
    y.sum().backward()
except RuntimeError:
    print("A saved value was modified")

Why exp needs its output

The derivative of exp(x) is exp(x), so its backward operation saves the forward output y. Assigning to y overwrites part of that saved result. For these ordinary autograd tensors, in-place operations update a version counter. When backward retrieves a saved tensor, it checks that the version has not changed.

The check is conservative: it protects the saved tensor even if you believe a particular overwritten entry will not affect the final gradient.

An in-place API is not a blanket safety guarantee

Some in-place operations support autograd in suitable contexts. That does not mean they are safe on every tensor. For example, an in-place ReLU may be valid on an intermediate, yet fail if it overwrites an output that an earlier operation saved. Leaf tensors requiring gradients have additional mutation restrictions.

no_grad prevents recording new operations; it does not restore overwritten values or disable version checks in an existing graph. A newly created tensor may already have been saved by its producing operation, as y was by exp. Neither rule makes arbitrary mutations safe.

Preserve the saved result

Construct the masked result separately:

x = torch.tensor([0., 1.], requires_grad=True)
y = x.exp()
mask = torch.tensor([True, False])
out = torch.where(mask, torch.zeros_like(y), y)
out.sum().backward()
print([round(v, 4) for v in out.tolist()])  # [0.0, 2.7183]
print([round(v, 4) for v in x.grad.tolist()])  # [0.0, 2.7183]

The original y remains available to exp's backward. At the replaced position, the output is a constant and its derivative is zero. Other positions retain the exponential derivative.

This example uses inputs with finite exponentials. Masking after an invalid or overflowing calculation does not generally make its gradients safe: where selects results but does not undo computation of its inputs.

def masked_exp(x: torch.Tensor, zero_at: int) -> torch.Tensor