class_docstr: Increase the learning rate of each parameter group from min lr to max lr
  over warmup_num_steps steps, and then fix at max lr.
  
  Args:
      optimizer (Optimizer): Wrapped optimizer.
      warmup_min_lr (float or list): minimum learning rate. Default: 0
      warmup_max_lr (float or list): maximum learning rate. Default: 0.001
      warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000
      warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log
      last_batch_iteration (int): The index of the last batch. Default: -1.
  Example:
      >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
      >>> scheduler = WarmupLR(optimizer)
      >>> data_loader = torch.utils.data.DataLoader(...)
      >>> for epoch in range(10):
      >>>     for batch in data_loader:
      >>>         train_batch(...)
      >>>         scheduler.step()
class_docstr_tok: ["Increase","the","learning","rate","of","each","parameter","group","from","min","lr","to","max","lr","over","warmup_num_steps","steps",",","and","then","fix","at","max","lr",".","Args",":","optimizer","(","Optimizer",")",":","Wrapped","optimizer",".","warmup_min_lr","(","float","or","list",")",":","minimum","learning","rate",".","Default",":","0","warmup_max_lr","(","float","or","list",")",":","maximum","learning","rate",".","Default",":","0.001","warmup_num_steps","(","int",")",":","number","of","steps","to","warm","up","from","min_lr","to","max_lr",".","Default",":","1000","warmup_type","{","\u2018","log","\u2019",",","\u2018","linear","\u2019","}",":","increasing","function","from","min_lr","to","max_lr","during","warmup",".","Default",":","log","last_batch_iteration","(","int",")",":","The","index","of","the","last","batch",".","Default",":","-1",".","Example",":",">",">",">","optimizer","=","torch.optim.SGD","(","model.parameters","(",")",",","lr=0.1",",","momentum=0.9",")",">",">",">","scheduler","=","WarmupLR","(","optimizer",")",">",">",">","data_loader","=","torch.utils.data.DataLoader","(","...",")",">",">",">","for","epoch","in","range","(","10",")",":",">",">",">","for","batch","in","data_loader",":",">",">",">","train_batch","(","...",")",">",">",">","scheduler.step","(",")"]
code: def step(self, last_batch_iteration=None):
      if last_batch_iteration is None:
          last_batch_iteration = self.last_batch_iteration + 1
      self.last_batch_iteration = last_batch_iteration
      for param_group, lr in zip(
          self.optimizer.param_groups, self.get_lr()
      ):
          param_group["lr"] = lr
      self._last_lr = [
          group["lr"] for group in self.optimizer.param_groups
      ]
  
imports: import argparse
  from torch.optim import Optimizer
  import math
  from deepspeed.utils import logger
label_desc: Use image node_id 4 for calling the WarmupLR obj's underlying member method code with example usage: obj.step(last_batch_iteration) without return types