Do you load all your data on the CUDA memory?
How much residual CUDA RAM do you have, after you instantiate your model?
Try with watch -n 1 nvidia-smi whilst running your code.
It’s rarely the model that uses all the memory, but either the copies for the batches, or the dataset itself IMHO, at least that is what’s happened to me many times.