def learning_rate_schedule(epoch):
    """Learning rate is scheduled to be reduced after 80 and 120 epochs.
    This function is automatically every epoch as part of callbacks
    during training.
    """

    if epoch < 80:
        return 1e-3
    if epoch < 120:
        return 1e-4
    return 1e-5