OOM error during validation in Transformer model

I am following the https://github.com/OpenNMT/OpenNMT-tf/tree/master/scripts/wmt tutorial to train a transformer for NMT. I get the OOM error from the GPU at the time of validation i.e. after every 1000 steps. I am providing a large training dataset and pre-trained embeddings.
I have tried lowering the training batch_size to 512 and the validation batch_size to 32.
I am using a GTX 1080 8 GB GPU.
How do we resolve this?

Thanks in advance


What is the size of your vocabulary?

I used the sentencepiece model to create the vocab file and my vocab size is 32000.