Just to confirm, run watch nvidia-smi. Ensure your GPU memory is near empty before u run the script and see what happens when u run it.
It could be that the model and data is in fact overloading the GPU, it happened to me when I tried running some 3D DL algo. Try batch_size=1, see how that goes.