Похоже, что глобальный шаг (tf.train.get_global_step) не был увеличен

Question

Я создаю EstimatorSpec как следующее:

optimizer = tf.keras.optimizers.Adam(0.0001)
checkpoint = tf.train.Checkpoint(
    step=tf.compat.v1.train.get_global_step(),
    optimizer=optimizer,
    net=model
)

# ...


return tf.estimator.EstimatorSpec(
    mode=mode,
    loss=loss,
    train_op=tf.group(
        optimizer.apply_gradients(zip(gradients, trainable_variables)),
        checkpoint.step.assign_add(1)
    ),
    eval_metric_ops=eval_metric_ops,
    evaluation_hooks=eval_hooks
)

Как вы можете видеть, я устанавливаю tf.compat.v1.train.get_global_step() на tf.train.Checkpoint и позже использую checkpoint.step.assign_add(1) при установке train_op для EstimatorSpec.

Моя тренировка работает, но в начале я получаю следующие предупреждения:

INFO:tensorflow:loss = 354.47638, step = 1
WARNING:tensorflow:It seems that global step (tf.train.get_global_step) has not been increased. Current value (could be stable): 1 vs previous value: 1. You could increase the global step by passing tf.train.get_global_step() to Optimizer.apply_gradients or Optimizer.minimize.
WARNING:tensorflow:It seems that global step (tf.train.get_global_step) has not been increased. Current value (could be stable): 4 vs previous value: 4. You could increase the global step by passing tf.train.get_global_step() to Optimizer.apply_gradients or Optimizer.minimize.
WARNING:tensorflow:It seems that global step (tf.train.get_global_step) has not been increased. Current value (could be stable): 19 vs previous value: 19. You could increase the global step by passing tf.train.get_global_step() to Optimizer.apply_gradients or Optimizer.minimize.

Это только для первых нескольких шагов. Я не знаю, почему они там в первую очередь - обучение работает - но кажется, что что-то не так, как должно быть.

Есть идеи?

Окружающая среда

Tensorflow 2.1.0