Good to hear you’ve found the issue!
Usually I would recommend to run your code on CPU and check if it fails, since the error messages are usually clearer than the CUDA ones.
Also, if you are running into an error using the GPU (while CPU works fine), you should launch your script using CUDA_LAUNCH_BLOCKING=1 python script.py args. This will synchronize your CUDA calls, such that the error should be thrown with the corresponding line of code. This cannor be ensured for the default behavior, since CUDA calls are asynchronous.