import tensorflow as tf


    output_spec = None
    if mode == tf.estimator.ModeKeys.TRAIN:
      train_op, update_learning_rate = optimization.create_optimizer(
          total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu, adjust_lr, use_hvd,
          use_compression, use_fp16, clip, cos_decay, use_lamb, previous_train_steps, post_train_steps)

      logging_hook = tf.train.LoggingTensorHook({"loss": total_loss, "learning_rate": update_learning_rate}, every_n_iter=FLAGS.hooking_frequence)
      output_spec = tf.estimator.EstimatorSpec(
          mode=mode,
          loss=total_loss,
          train_op=train_op,
          training_hooks=[logging_hook])
    elif mode == tf.estimator.ModeKeys.EVAL:
