Thanks for your reply. But this seems to only work on layer level?
Currently, I got weird nan error in backward pass. I need to investigate the gradient of every variable to see which part goes wrong.
Idealy something similar to the following:
def inves(name=''):
    def f(tensor):
        if np.isnan(torch.mean(tensor).data.cpu().numpy() ):
            print('gradient of {} is'.format(name))
            print(tensor)
            assert 0, 'nan gradient'
    return f

for key, var in Model.varibles().iteritems():
    var.register_hook(inves(key))