Tensorflow2 seq2seq neural machine translation model - excessive memory copy between host and device

Question

I am training the seq2seq nmt model provided as an official tutorial https://colab.research.google.com/github/tensorflow/addons/blob/master/docs/tutorials/networks_seq2seq_n