Запустить тренировку нейронной сети на Google Colab TPU

Есть модель нейронной сети Embedding GRU для генерации текста.
Код взят отсюда https://www.tensorflow.org/text/tutorials/text_gen...
Необходимо настроить параметры датасета и запустить тренировку модели на Google Colab TPU.
Текущий код:
https://colab.research.google.com/drive/1FCOamhAX1...
Все остальные материалы отправлю потенциальному исполнителю.