Error in saving a Neural machine translation tutorials model

Hi.
I tried to save a tutorial model: https://www.tensorflow.org/text/tutorials/nmt_with_attention
after training it but i got an error: ‘DatasetSpec’ object has no attribute ‘shape’.
what i should do to resolve this issue?

Are you using tf.saved_model.save to save the model ?