Solved by Varal7 in post #2 
              
              
                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

Howe…