как распечатать скорость обучения для каждой эпохи, в которой используется функциональный планировщик скорости обучения?

Я использовал настраиваемый планировщик скорости обучения. Код выглядит следующим образом (То же, что и https://www.tensorflow.org/tutorials/text/transformer#optimizer).

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
   def __init__(self, d_model, warmup_steps=4000):
     super(CustomSchedule, self).__init__()

     self.d_model = d_model
     self.d_model = tf.cast(self.d_model, tf.float32)

     self.warmup_steps = warmup_steps

  def __call__(self, step):
    arg1 = tf.math.rsqrt(step)
    arg2 = step * (self.warmup_steps ** -1.5)

    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)

Я использую model.fit. Хотите печатать lr в конце каждой эпохи.