I think the canonical reference for finding bad gradients is this snippet by Adam Paszke:

  
      gist.github.com
  
  
    https://gist.github.com/apaszke/f93a377244be9bfcb96d3547b9bc424d 2.5k
bad_grad_viz.py
from graphviz import Digraph
import torch
from torch.autograd import Variable, Function

def iter_graph(root, callback):
    queue = [root]
    seen = set()
    while queue:
        fn = queue.pop()
        if fn in seen:
This file has been truncated. show original




  
  
    
    
  
  


It checks for NaN (by using x!=x if and only if x is NaN) and very large gradients, but you could easily adapt is_bad_grad to best fit your purpose.
Best regards
Thomas