def METHOD_NAME(self, step_num=None):
    """Set lr for all parameters groups."""
    if step_num is None:
        step_num = self.num_iters + 1
    self.num_iters = step_num
    new_lr = self.get_lr()
    for group in self.optimizer.param_groups:
        group["lr"] = new_lr