def _scheduler_step(scheduler, reduced_metric=None):
if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
scheduler.step(reduced_metric)
else:
scheduler.step()
lr_list = [param_group["lr"] for param_group in scheduler.optimizer.param_groups]
momentum_list = get_optimizer_momentum_list(scheduler.optimizer)
return lr_list, momentum_list
def _get_lr_momentum_stats(self) -> Dict:
lr_list = [param_group["lr"] for param_group in self.optimizer.param_groups]
momentum_list = get_optimizer_momentum_list(self.optimizer)
stats = {self._prefix_lr: lr_list[0], self._prefix_momentum: momentum_list[0]}
return stats