class_docstr: Tristage learning rate schedulr
  
  Implement the learning rate scheduler in https://arxiv.org/pdf/1904.08779.pdf
  
  Similar to inverse_squre_root scheduler, but tri_stage learning rate employs
  three stages LR scheduling:
  
      - 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 = cfg.init_lr_scale * cfg.lr
    lrs = torch.linspace(init_lr, cfg.lr, cfg.warmup_steps)
    lr = lrs[update_num]
  
  During hold::
  
    lr = cfg.lr
  
  During decay::
  
    decay_factor = - math.log(cfg.final_lr_scale) / cfg.decay_steps
    lr = cfg.lr * exp(- (update_num - warmup_steps - decay_steps) * decay_factor)
  
  After that::
  
    lr = cfg.lr * cfg.final_lr_scale
docstr: Update the learning rate after each update.
docstr_tok: ["Update","the","learning","rate","after","each","update","."]
code: def step_update(self, num_updates):
      """Update the learning rate after each update."""
      stage, steps_in_stage = self._decide_stage(num_updates)
      if stage == 0:
          self.lr = self.init_lr + self.warmup_rate * steps_in_stage
      elif stage == 1:
          self.lr = self.peak_lr
      elif stage == 2:
          self.lr = self.peak_lr * math.exp(
              -self.decay_factor * steps_in_stage
          )
      elif stage == 3:
          self.lr = self.final_lr
      else:
          raise ValueError("Undefined stage")
  
      self.optimizer.set_lr(self.lr)
  
      return self.lr
  
file_path: fairseq/fairseq/optim/lr_scheduler/tri_stage_lr_scheduler.py
imports: import math
  from dataclasses import dataclass, field
  from typing import Optional, List, Tuple
  from omegaconf import II
  from fairseq.dataclass import FairseqDataclass
  from fairseq.optim.lr_scheduler import FairseqLRScheduler, register_lr_scheduler
label_desc: Use image node_id 4 for calling the TriStageLRSchedule obj's underlying member method code with example usage: obj.step_update(num_updates) and returns: self