CUDA out of memory. Tried to allocate 314.00 MiB. GPU Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace

Yes when i decrease the batch size to 2 it worked thanks