Solved by JuanFMontesinos in post #2 
              
              
                The gradient is input.grad 
What do you understand by loss.grad? 
input.grad is gradient of loss wrt input which is the cross entropy gradient. 
there is no loss.grad as it is not involved in further opts 
import torch

input = torch.randn(3, 5, requires_grad=True)
input.register_hook(lambda x: prin…