LeoSerena:

Just an auxiliary question though: do we agree that this does not update the .grad value of the layer and that I can still call .backward() like normal afterwards?


Yes, as seen here:
lin = nn.Linear(1, 1)
x = torch.randn(1, 1)

out = lin(x)

grads = torch.autograd.grad(out, lin.parameters())
print(grads)
> (tensor([[1.3361]]), tensor([1.]))
print([p.grad for p in lin.parameters()])
> [None, None]