Ошибка TPU Colab InvalidArgumentError: Невозможно назначить устройство для работы

Question

в Google Colab при использовании TPU, у меня появляется следующая ошибка

InvalidArgumentError: Невозможно назначить устройство для операции Adam / iterations / IsInitialized / VarIsInitializedOp: {{узел Adam / iterations / IsInitialized / VarIsInitializedOp}} былявно назначены на / job: worker / replica: 0 / task: 0 / device: TPU: 0, но доступны следующие устройства [/ job: localhost / replica: 0 / task: 0 / device: CPU: 0, / job: localhost /реплика: 0 / задача: 0 / устройство: XLA_CPU: 0]. Убедитесь, что в спецификации устройства указано действительное устройство. [[Адам / итерации / IsInitialized / VarIsInitializedOp]]

TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
        resolver = tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)
        tf.contrib.distribute.initialize_tpu_system(resolver)
        strategy = tf.contrib.distribute.TPUStrategy(resolver)
        with strategy.scope():
            # Setup the model inputs / outputs
            model = Model(inputs=[inputs_img, inputs_mask], outputs=outputs)
            # Compile the model
            model.compile(
                optimizer = Adam(lr=lr),
                loss=self.loss_total(inputs_mask)
            )