def lr_decay(N, step, learning_rate):
    """
    learning rate decay
    Args:
        learning_rate: base learning rate
        step: current iteration number
        N: total number of iterations over which learning rate is decayed
    """
    min_lr = 0.00001
    res = learning_rate * ((N - step) / N) ** 2
    return max(res, min_lr)