CUDA operations are executed asynchronously, so a failed operation might raise an assertion, which will be triggered by the next CUDA operation.
In your case an indexing failed and you see the error in the next .to() operation.
You can rerun the code via:
CUDA_LAUNCH_BLOCKING=1 python script.py args

to get the proper line of code which is failing (or run the code on the CPU alternatively).