Ошибка загрузки предварительно обученного BERT на ТПУ с использованием Keras

Question

Я работаю с TensorFlow и Keras на TPU, и мне нужно загрузить предварительно обученную модель BERT и преобразовать ее в модель TPU. Я сначала загружаю модель, и она отлично работает.

def _load_bert(self):

    logging.info('Loading BERT from %s', self.bert_path)
    config_path = os.path.join(self.bert_path, 'bert_config.json')
    checkpoint_path = os.path.join(self.bert_path, 'bert_model.ckpt')

    bert = keras_bert.load_trained_model_from_checkpoint(
        config_path,
        checkpoint_path
    )

    self._bert = bert
    self._seq_length = int(self._bert.input[0].shape[1])

А потом я пытаюсь преобразовать его в модель TPU, выполнив следующее:

self._bert = tensorflow.contrib.tpu.keras_to_tpu_model(
    self._bert,
    strategy=tensorflow.contrib.tpu.TPUDistributionStrategy(
        tensorflow.contrib.cluster_resolver.TPUClusterResolver(
            self.tpu, zone=self.tpu_zone, project=self.gcp_project
        )
    )
)

А потом почему-то я получил следующую ошибку:

ValueError: ('Expected `model` argument to be a `Model` instance, got ', <keras.engine.training.Model object at xxxxxxxxxx >)

Это работает, когда я не конвертирую его в модель TPU. У кого-нибудь есть подсказка?