def get_inverse_square_root_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, warmup_init_lr=0.0, last_epoch=-1
):
    """
    Create a schedule with a learning rate that Decay the LR based on the inverse square root of the update number.
    After a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        warmup_init_lr (:obj:`float`, `optional`, defaults to 0):
            The initial LR for warmup.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    
    During warmup::
      lrs = torch.linspace(arg.warmup_init_lr, arg.lr, arg.warmup_updates)
      lr = lrs[update_num]
    After warmup::
      decay_factor = arg.lr * sqrt(arg.warmup_updates)
      lr = decay_factor / sqrt(update_num)

    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    lr = optimizer.defaults["lr"]
    assert lr > warmup_init_lr, f"lr ({lr}) must be be bigger than initial lr ({warmup_init_lr})"

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            lr_step = (lr - warmup_init_lr) / num_warmup_steps
            return (warmup_init_lr + current_step * lr_step) / lr
        elif current_step > num_training_steps:
            return 1e-7 / lr  # as LambdaLR multiplies by lr_init
        else:
            decay_factor = lr * num_warmup_steps**0.5
            return (decay_factor * current_step**-0.5) / lr

    return LambdaLR(optimizer, lr_lambda, last_epoch)

    Create a schedule with a learning rate that Decay the LR based on the inverse square root of the update number.
    After a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        warmup_init_lr (:obj:`float`, `optional`, defaults to 0):
            The initial LR for warmup.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    
    During warmup::
      lrs = torch.linspace(arg.warmup_init_lr, arg.lr, arg.warmup_updates)
      lr = lrs[update_num]
    After warmup::
      decay_factor = arg.lr * sqrt(arg.warmup_updates)
      lr = decay_factor / sqrt(update_num)

    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    
uer/utils/optimizers.py