Hi,
Could you enable anomaly mode to see where it comes from? torch.autograd.set_detect_anomaly(True) at the begining of your script.
Also if you have custom autograd.Function, could you share their code here?