code: def set_lr(self, lr):
          
          if isinstance(lr, list):
              for group_lr, param_group in zip(lr, self.optimizer.param_groups):
                  param_group['lr'] = group_lr
          else:
              for param_group in self.optimizer.param_groups:
                  param_group['lr'] = lr
docs: Set a learning rate for the optimizer