Hi,
Are you sharing the GPU with other people?
Pytorch reports <3GB used but cuda report 11+GB in use on the device ( in the case where it fails).
As you can see, in the case where it works with a bigger batch size, pytorch actually uses much more memory: 5/6GB.
If there is a lot of memory already used on the gpu even though nothing is running (you can check with nvidia-smi). You can use (sudo) lsof /dev/nvidiaX where X is the number of the gpu you want to check. This will list all process using the gpu more thoroughly than nvidia-smi. If there are some things that shouldn’t be here like old python process, you can get rid of them by killing them.