code: def optimize_fn(model,
                  optimizer=None,
                  learning_rate_schedule=None,
                  loss=None,
                  metrics=None):
    
    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)
docs: Compile the model in Keras.