Hi Antoine,
You can use torch.set_grad_enabled(True) to activate gradient tracking even under a no_grad context.
For instance:
import torch

x = torch.ones(5, requires_grad=True)
with torch.no_grad():
  torch.set_grad_enabled(True)
  y = x * x

y.sum().backward()
print(x.grad)  # works fine

However, this does not work if you disabled gradients with the torch.inference_mode context:
x = torch.ones(5, requires_grad=True)
with torch.inference_mode():
  torch.set_grad_enabled(True)
  y = x * x

y.sum().backward()  # fails