Hi,
In theory, that is what the register_backward_hook() on nn.Module is supposed to give you.
Unfortunately, they are broken right now and we will hopefully be able to fix them later this year.
The way you can do this right now is to use register_forward_hook() on your nn.Module to access both the input and outputs and use register_hook() on these Tensors to access their gradients.