docs: def get_train_op(loss, params):
    """Generate training operation that updates variables based on loss."""
    with tf.variable_scope("get_train_op"):
      mlperf_log.transformer_print(
          key=mlperf_log.OPT_LR_WARMUP_STEPS,
          value=params.learning_rate_warmup_steps)
      learning_rate = get_learning_rate(
          params.learning_rate, params.hidden_size,
          params.learning_rate_warmup_steps)
      log_id = mlperf_log.resnet_print(key=mlperf_log.OPT_LR, deferred=True)
      learning_rate = tf_mlperf_log.log_deferred(op=learning_rate, log_id=log_id,
                                                 every_n=100)
  
      # Create optimizer. Use LazyAdamOptimizer from TF contrib, which is faster
      # than the TF core Adam optimizer.
      mlperf_log.transformer_print(key=mlperf_log.OPT_NAME,
                                   value=mlperf_log.LAZY_ADAM)
      mlperf_log.transformer_print(key=mlperf_log.OPT_HP_ADAM_BETA1,
                                   value=params.optimizer_adam_beta1)
      mlperf_log.transformer_print(key=mlperf_log.OPT_HP_ADAM_BETA2,
                                   value=params.optimizer_adam_beta2)
      mlperf_log.transformer_print(key=mlperf_log.OPT_HP_ADAM_EPSILON,
                                   value=params.optimizer_adam_epsilon)
      optimizer = tf.contrib.opt.LazyAdamOptimizer(
          learning_rate,
          beta1=params.optimizer_adam_beta1,
          beta2=params.optimizer_adam_beta2,
          epsilon=params.optimizer_adam_epsilon)
  
      # Calculate and apply gradients using LazyAdamOptimizer.
      global_step = tf.train.get_global_step()
      tvars = tf.trainable_variables()
      gradients = optimizer.compute_gradients(
          loss, tvars, colocate_gradients_with_ops=True)
      train_op = optimizer.apply_gradients(
          gradients, global_step=global_step, name="train")
  
      # Save gradient norm to Tensorboard
      tf.summary.scalar("global_norm/gradient_norm",
                        tf.global_norm(list(zip(*gradients))[0]))
  
      return train_op