Could you rerun the script with:
CUDA_LAUNCH_BLOCKING=1 python script.py args

and post the stack trace here, please?
The illegal memory access might have been created by a previous CUDA operation and your loss could be a red herring.