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 step(self, closure: Any = None) -> None:
      super().step(closure)
      iter_, stage_id = self._get_warmup_state()
  
      iter_ += 1
      if iter_ > self._stages[
          stage_id
      ].max_iters and stage_id + 1 < len(self._stages):
          stage_id += 1
          logger.info(
              "Warmup Optimizer finishing "
              f"{self._stages[stage_id - 1]} "
              "switching to "
              f"{self._stages[stage_id]}"
          )
      self._set_lr(iter_, stage_id)
  
      # pyre-ignore [16]
      self.state[self._warmup_param] = {
          "warmup": torch.tensor([iter_, stage_id], dtype=torch.long)
      }
  
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 5 for calling the WarmupOptimizer obj's underlying member method code with example usage: obj.step(closure) without return types