class_docstr: Adjusts learning rate according to the schedule.
  
  Args:
      optimizer (KeyedOptimizer): optimizer to wrap
      stages (List[WarmupStage]): stages to go through
      lr (float): initial learning rate
      lr_param (str): learning rate parameter in parameter group.
      param_name: Name of fake parameter to hold warmup state.
class_docstr_tok: ["Adjusts","learning","rate","according","to","the","schedule",".","Args",":","optimizer","(","KeyedOptimizer",")",":","optimizer","to","wrap","stages","(","List","[","WarmupStage","]",")",":","stages","to","go","through","lr","(","float",")",":","initial","learning","rate","lr_param","(","str",")",":","learning","rate","parameter","in","parameter","group",".","param_name",":","Name","of","fake","parameter","to","hold","warmup","state","."]
code: def __init__(
      self,
      optimizer: KeyedOptimizer,
      stages: List[WarmupStage],
      lr: float = 0.1,
      lr_param: str = "lr",
      param_name: str = "__warmup",
  ) -> None:
      super().__init__(optimizer)
      self._stages: List[WarmupStage] = _lr_stages(stages)
      self._lr_param: str = lr_param
      self._lr: float = lr
      self._warmup_param: torch.nn.Parameter = torch.nn.Parameter()
      # pyre-ignore [16]
      self.params[param_name] = self._warmup_param
      # for fused optimizer we will do first backward() pass before calling step()
      self._set_lr(0, 0)
  
file_path: torchrec/torchrec/optim/warmup.py
imports: import logging
  import math
  from dataclasses import dataclass
  from enum import Enum, unique
  from typing import Any, List, Tuple
  import torch
  from torchrec.optim.keyed import KeyedOptimizer, OptimizerWrapper
label_desc: Use image node_id 1 to create a new WarmupOptimizer object from inherited base classes: OptimizerWrapper with example: obj = WarmupOptimizer(optimizer, stages, lr, lr_param, param_name)