Missing get_last_lr() in Custom PolyLR Scheduler

[SimoneBendazzoli93]

The custom PolyLRScheduler scheduler implementation is missing the get_last_lr() function, which is essential for retrieving the last computed learning rate for each parameter group. This function is present in PyTorch’s built-in schedulers and is useful for logging, debugging, and maintaining API consistency. Adding this function will improve usability and integration with standard PyTorch workflows.
One practical example is integrating nnUNet training in a MONAI Bundle standard (i.e. using PolyLRScheduler in the LearningRateSchedulerHandler ): currently, in order to do so, the PolyLRScheduler needs to be modified as follows:

from torch.optim.lr_scheduler import _LRScheduler


class PolyLRScheduler(_LRScheduler):
    def __init__(self, optimizer, initial_lr: float, max_steps: int, exponent: float = 0.9, current_step: int = None):
        self.optimizer = optimizer
        self.initial_lr = initial_lr
        self.max_steps = max_steps
        self.exponent = exponent
        self.ctr = 0
        super().__init__(optimizer, current_step if current_step is not None else -1, False)

    def step(self, current_step=None):
        if current_step is None or current_step == -1:
            current_step = self.ctr
            self.ctr += 1

        new_lr = self.initial_lr * (1 - current_step / self.max_steps) ** self.exponent
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = new_lr

        self._last_lr = [group['lr'] for group in self.optimizer.param_groups] # 🚀 Missing definition

    def get_last_lr(self):   # 🚀 Missing function that should be included
        return self._last_lr