You could try that:
def register_nan_checks(model):
    def check_grad(module, grad_input, grad_output):
        if np.isnan(grad_input.data.numpy()):
            print('NaN gradient in ' + type(module).__name__)
    model.apply(lambda module: module.register_backward_hook(check_grad))