def get_lr_warmup_schedule(lr_final, schedule_end_step, schedule_start_step=0, lr_init=1e-10):
    """Creates a learning rate warm-up schedule.

    The learning rate is increased linearly during the steps.

    Arguments:
        lr_final: A float representing the value of the learning rate at the end of the warmup.
        schedule_start_step: An integer representing the step (either iteration or epoch) after
            which to start applying the warmupi.
        schedule_end_step: An integer representing the step (either iteration or epoch) after
            which we stop applying the warmup, and the learning rate reaches its final value
            `lr_final`.
         lr_init: A float representing the value of the learning rate at the beginning of the
            warm-up.
    Returns
        lr: A float representing the learning rate.
    """
    total_steps = schedule_end_step - schedule_start_step
    total_lr_diff = lr_final - lr_init

    def lr_schedule(step, lr):
        if schedule_start_step <= step < schedule_end_step:
            lr = lr_init + total_lr_diff * step / total_steps
            # print('Step ', step, '      lr=', lr)

        return lr

    return lr_schedule