Thanks @tom
But when I use it in my network:
...
import bad_grad_viz as bad_grad
...
loss_G , loss_D = model(data_1, data_2)
get_dot = bad_grad.register_hooks(loss_1)
model.optimizer_network.zero_grad()
loss_G.backward()
loss_D.backward()
dot = get_dot()
dot.save('/path/to/dir/tmp.dot')

I get this error from bad_grad_viz.py 175:
grad_output = grad_output.data
AttributeError: 'NoneType' object had no attribute 'data'

More info:
model is a GAN network.
print(type(loss_G))
>>> <class 'torch.autograd.variable.Variable'>
print(type(loss_D))
>>> <class 'torch.autograd.variable.Variable'>
print(loss_G.requires_grad)
>>> True
print(loss_D.requires_grad)
>>> True

Any suggestion @apaszke ?