I want to save gradients of internal variables through register_hook() or retain_grad().
When I run model on single GPU, it works.
But when I run model on multiple GPUs through wrapping model into nn.DataParallel, I find that it doesn’t work.
Can anyone help me?