it’s possible that moving both the training and test dataset over to CUDA doesn’t give enough space for the network to forward-prop through. Loading your dataset into GPU memory itself takes 1.3 GB of memory.
Are you giving a large batch size as input to your network? Can you reduce the batch size at inference time?