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 at the end of the given epoch.
docstr_tok: ["Update","the","learning","rate","at","the","end","of","the","given","epoch","."]
code: def step(self, epoch, val_loss=None):
      """Update the learning rate at the end of the given epoch."""
      super().step(epoch, val_loss)
      # we don't change the learning rate at epoch boundaries
      return self.optimizer.get_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 3 for calling the TriStageLRSchedule obj's underlying member method code with example usage: obj.step(epoch, val_loss) and returns: self