def METHOD_NAME(self, increment=1):
    """Set lr for all parameters groups."""
    self.num_steps += increment
    new_lr = self.get_lr()
    for group in self.optimizer.param_groups:
        group['lr'] = new_lr