Hi,
One thing you can do is check the line numbers in the traceback. If you look up the class the line belongs to in torch/autograd/_functions/*.py, you can tell which operation it bails out on to narrow it down.
I wonder whether it might be worth adding a “debug” modus that records the stack of the op in the forward pass and spits it out on error in the backward. That way, it would point to the right line of code directly.
Best regards
Thomas