Не получается сохранить модель tensorflow

Есть рабочий код для создания чат бота на архитектуре transformer. Использую библиотеку tensorflow. Хочу сохранить модель:

model.save('model/mm.h5')

Получаю ошибку:

TypeError: Unable to serialize 256.0 to JSON. Unrecognized type <class 'tensorflow.python.framework.ops.EagerTensor'>.

Полный код:

https://github.com/Mik-404/tf2-transformer-chatbot/blob/master/tf2_tpu_transformer_chatbot.ipynb