If you need to have more flexibility, you could have a look at forward hooks in this post.
Since you need the gradients, make sure not to detach
the tensor.
1 Like
If you need to have more flexibility, you could have a look at forward hooks in this post.
Since you need the gradients, make sure not to detach
the tensor.