Hi,

Yes using .detach() is the right way to stop gradients from flowing back.

To get the gradients, you don’t have to make a full backward, you can use torch.autograd.grad 49 to get the gradients for specific tensors, here embedding for example. You will need to give create_graph=True (which implies will set retain_graph=True).
Also I am not sure why you set training=False on the new net you create?
The last .backward() in the update function should not have a retain_graph=True. If you get an error here, that means that accross two calls to update, some autograd operations are shared (and they shouldn’t).