def lr_scheduler(epoch, init_lr=0.001, lr_decay_epoch=7):
    """
    Returns the current learning rate given the epoch. This decays the learning rate
    by a factor of 0.1 every lr_decay_epoch epochs.
    """
    return init_lr * (0.1**(epoch // lr_decay_epoch))