PyTorch `torch.no_grad` vs `torch.inference_mode`

autograd
2

Yes, you can depend on runtime errors and as long as no errors are raised, you code should be fine.
One difference would be that you are not allowed to set the requires_grad attribute on tensors from an inference_mode context:

with torch.no_grad():
    x = torch.randn(1)
    y = x + 1

y.requires_grad = True
z = y + 1
print(z.grad_fn)
> <AddBackward0 object at 0x7fe9c6eafdf0>

with torch.inference_mode():
    x = torch.randn(1)
    y = x + 1

y.requires_grad = True
> RuntimeError: Setting requires_grad=True on inference tensor outside InferenceMode is not allowed.
21 Likes