Hi,
Yes I think it comes from there.
Does this error happen during the forward or the backward pass?

If forward, can you give the stack trace that shows which function is faulty?
If backward, you can enable the anomaly mode (torch.autograd.set_detect_anomaly(True)) to get a warning that will tell you which forward function is responsible for the error in backward. And you can send it here.