Точная настройка BERT с помощью оценщиков на TPU на colab TypeError: неподдерживаемые типы операндов для * =: 'NoneType' и 'int'

Question

я написал блокнот jupyter на colab google для точной настройки (для классификации текста) версии BERT, которую я уже предварительно обучил только на арабском языке.И я не смог обойти эту ошибку, когда началось обучение.

Я последовал за записной книжкой, которую дал Google на github

Код модели здания:

model_fn = model_fn_builder(
  bert_config=modeling.BertConfig.from_json_file(CONFIG_FILE),
  num_labels=len(label_list),
  init_checkpoint=INIT_CHECKPOINT,
  learning_rate=LEARNING_RATE,
  num_train_steps=num_train_steps,
  num_warmup_steps=num_warmup_steps,
  use_tpu=True,
  use_one_hot_embeddings=True
)


tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)

run_config = tf.contrib.tpu.RunConfig(
    cluster=tpu_cluster_resolver,
    model_dir=OUTPUT_DIR,
    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS,
    tpu_config=tf.contrib.tpu.TPUConfig(
        iterations_per_loop=ITERATIONS_PER_LOOP,
        num_shards=NUM_TPU_CORES,
        per_host_input_for_training=tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2))

estimator = tf.contrib.tpu.TPUEstimator(
    use_tpu=USE_TPU,
    model_fn=model_fn,
    config=run_config,
    train_batch_size=TRAIN_BATCH_SIZE,
    eval_batch_size=EVAL_BATCH_SIZE,
    predict_batch_size=PREDICT_BATCH_SIZE,)

train_input_fn = input_fn_builder(
    features=train_features,
    seq_length=MAX_SEQ_LENGTH,
    is_training=True,
    drop_remainder=False)

#tf.reset_default_graph()
print(f'Beginning Training!')
current_time = datetime.now()
estimator.train(input_fn=train_input_fn, max_steps=TRAIN_STEPS)
print("Training took time ", datetime.now() - current_time)

код ошибки:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/tpu/tpu_sharding.py in _unshard_shape(self, shape)
    214                        (shape.as_list(), self._shard_dimension))
    215     dims = shape.as_list()
--> 216     dims[self._shard_dimension] *= self._number_of_shards
    217     return tensor_shape.as_shape(dims)
    218 

TypeError: unsupported operand type(s) for *=: 'NoneType' and 'int'

Параметры и остальной код находятся в этой общей копии записной книжки colab: colab_link