Модель обучения на Colab TPU с распределенной стратегией

Я пытаюсь обучить и запустить модель классификации изображений в Colab, используя TPU. Нет питорча.

Я знаю, что TPU работает только с файлами из корзин GCS, поэтому я загружаю набор данных из корзины, а также прокомментировал функции контрольной точки и ведения журнала, чтобы не было ошибок такого типа. Я просто хочу, чтобы он тренировался без ошибок на ТПУ.

На CPU и GPU код работает, но проблема возникает, когда я использую with strategy.scope(): перед созданием модели. Это функция, которая доставляет мне проблемы при обучении модели:

def train_step(self, images, labels):
    with tf.GradientTape() as tape:
        predictionProbs = self(images, training=True)
        loss = self.loss_fn(labels, predictionProbs)

    grads = tape.gradient(loss, self.trainable_weights)

    predictionLabels = tf.squeeze(tf.cast(predictionProbs > PROB_THRESHOLD_POSITIVE, tf.float32), axis=1)
    acc = tf.reduce_mean(tf.cast(predictionLabels == labels, tf.float32))

    self.optimizer.apply_gradients(zip(grads, self.trainable_weights)) # here is the problem

    return loss, acc

И это ошибка, с которой я сталкиваюсь:

RuntimeError: `apply_gradients() cannot be called in cross-replica context. Use `tf.distribute.Strategy.run` to enter replica context.

Я посмотрел на https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy и я думаю, что вот решение, но я никогда не делал этого раньше, и я не знаю, откуда Я могу начать.

Может кто-нибудь, пожалуйста, дать мне совет по этой проблеме?