def METHOD_NAME(epoch):
    # Adjust epoch to produce round numbers at the end of each epoch, so that TensorBoard
    # learning rate graphs look better.
    epoch += 1.0 / self.steps_per_epoch
    return 1.0 / n_replicas * (epoch * (n_replicas - 1) / warmup_epochs + 1)