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.