Hi @ptrblck, I was actually able to solve the issue after a couple days of debugging. When I did run the script with CUDA_LAUNCH_BLOCKING=1 it actually just returned an identical traceback and error, minus the tip telling me to run with the flag, so I just had to change things until it worked
The issue ended up being that even though I created the accumulated loss tensor as a 0D tensor, when I accumulated loss over each sample using torch.add I was unaware that torch.cond sometimes returns a 0D tensor, sometimes a 1D tensor with 1 element depending on which method you select, which you can see in their examples section. This generated some unexpected behavior since I switched the method of calculating the condition number around the same time I added the exception handling, so I was calling backward on a tensor of size [1] when using my accumulated loss
A couple things seem strange to me: 1) why would I get such an ambiguous error message for a simple error of calling backward on a tensor of the wrong dimension? and 2) torch.cond returns tensors of different dimensions for different methods without much mention in the documentation
The latter is what burned me I think. All I have to do is call torch.squeeze on the output of torch.cond before accumulating the loss and now it works fine