def schedule(epoch, learning_rate):
    if epoch < 2:
        return learning_rate
    return learning_rate * 0.5