Generate training operation that updates variables based on loss.

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