Compile the model in Keras.

def optimize_fn(model,
                optimizer=None,
                learning_rate_schedule=None,
                loss=None,
                metrics=None):
  """Compile the model in Keras."""
  learning_rate_schedule = learning_rate_schedule or T2TLearningRateSchedule()
  if optimizer:
    optimizer = optimizer(learning_rate=learning_rate_schedule)
  else:  # We use Adam by default with adjusted parameters.
    optimizer = tf.keras.optimizers.Adam(
        learning_rate=learning_rate_schedule,
        beta_1=0.9, beta_2=0.997, epsilon=1e-9)
  metrics = metrics or [tf.keras.metrics.sparse_categorical_accuracy]
  def xent_loss(y, x):
    return tf.keras.backend.sparse_categorical_crossentropy(
        y, x, from_logits=True)
  loss = loss or xent_loss
  return model.compile(optimizer=optimizer,
                       loss=loss,
                       metrics=metrics)