Solved by albanD in post #2 
              
              
                Hi, 
The hook does not take the tensor as input but the gradient. 
So the right signature would be: 
def hook_fn(self, grad):
    pass

So you directly have access to the gradient and can do whatever you want.