def get_tri_stage_schedule(optimizer, num_warmup_steps, num_decay_steps, num_training_steps, init_lr_scale=0.01, final_lr_scale=0.05, last_epoch=-1):
    """
    Create a schedule with a learning rate that have three stages: a warmup stage, a hold stage and a decay stage.
    Implement the learning rate scheduler in https://arxiv.org/pdf/1904.08779.pdf

        - warmup stage, starting from `lr` * `init_lr_scale`, linearly
          increased to `lr` in `warmup_steps` iterations
        - hold stage, after `warmup_steps`, keep the LR as `lr` for `hold_steps`
          iterations
        - decay stage, after hold stage, decay LR exponetially to
          `lr` * `final_lr_scale` in `decay_steps`;
          after that LR is keep as `final_lr_scale` * `lr`

    During warmup::
      init_lr = arg.init_lr_scale * arg.lr
      lrs = torch.linspace(init_lr, arg.lr, arg.warmup_steps)
      lr = lrs[update_num]
    During hold::
      lr = arg.lr
    During decay::
      decay_factor = - math.log(arg.final_lr_scale) / arg.decay_steps
      lr = arg.lr * exp(- (update_num - warmup_steps - decay_steps) * decay_factor)
    After that::
      lr = arg.lr * arg.final_lr_scale

    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_decay_steps (:obj:`int`):
            The number of steps for the decay phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        decay_scale (:obj:`float`):

        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """
    lr_hold = optimizer.defaults["lr"]
    lr_int = lr_hold * init_lr_scale
    lr_end = lr_hold * final_lr_scale

    def lr_lambda(current_step: int):
        warmup_rate = (lr_hold - lr_int) / num_warmup_steps
        decay_factor = -math.log(final_lr_scale) / max(num_decay_steps, 1)

        if current_step < num_warmup_steps:
            return (lr_int + current_step * warmup_rate) / lr_hold
        elif current_step >= num_warmup_steps and current_step < num_training_steps - num_decay_steps:
            return 1
        elif current_step <= num_training_steps:
            return math.exp(-decay_factor * (current_step - num_training_steps + num_decay_steps))
        else:
            return lr_end / lr_hold

    return LambdaLR(optimizer, lr_lambda, last_epoch)

    Create a schedule with a learning rate that have three stages: a warmup stage, a hold stage and a decay stage.
    Implement the learning rate scheduler in https://arxiv.org/pdf/1904.08779.pdf

        - warmup stage, starting from `lr` * `init_lr_scale`, linearly
          increased to `lr` in `warmup_steps` iterations
        - hold stage, after `warmup_steps`, keep the LR as `lr` for `hold_steps`
          iterations
        - decay stage, after hold stage, decay LR exponetially to
          `lr` * `final_lr_scale` in `decay_steps`;
          after that LR is keep as `final_lr_scale` * `lr`

    During warmup::
      init_lr = arg.init_lr_scale * arg.lr
      lrs = torch.linspace(init_lr, arg.lr, arg.warmup_steps)
      lr = lrs[update_num]
    During hold::
      lr = arg.lr
    During decay::
      decay_factor = - math.log(arg.final_lr_scale) / arg.decay_steps
      lr = arg.lr * exp(- (update_num - warmup_steps - decay_steps) * decay_factor)
    After that::
      lr = arg.lr * arg.final_lr_scale

    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_decay_steps (:obj:`int`):
            The number of steps for the decay phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        decay_scale (:obj:`float`):

        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    
uer/utils/optimizers.py