def constant_schedule_with_warmup(epoch, warmup_epochs=0, lr_start=1e-4, lr_max=1e-3):
    """ Create a schedule with a constant learning rate preceded by a warmup
    period during which the learning rate increases linearly between {lr_start} and {lr_max}.
    """

    if epoch < warmup_epochs:
        lr = (lr_max - lr_start) / warmup_epochs * epoch + lr_start
    else:
        lr = lr_max

    return lr