You should just add at the beginning or your program torch.autograd.set_detect_anomaly(True).
Then you will see a second stack trace appear above the one in the error. And that one will point you to the forward call corresponding to the backward error and that will help you identify which part of your model is failing