Hi,
You can try to enable torch.autograd.set_detect_anomaly(True) to see which forward function corresponds to the failing one in the backward.
Also You can try to run the same code on CPU to see if you get a better error message.
Also, I would double check all the indices in the multiple indexing you do.
Finally. You should never use .data anymore. This will reduce the chance of weird bugs a lot