def METHOD_NAME(self) -> List[float]:
    """Called when epoch is passed as a param to the `step` function of the scheduler."""
    if self.last_epoch < self.warmup_epochs:
        return [
            self.warmup_start_lr
            + self.last_epoch * (base_lr - self.warmup_start_lr) / max(1, self.warmup_epochs - 1)
            for base_lr in self.base_lrs
        ]
    return [
        self.eta_min
        + 0.5
        * (base_lr - self.eta_min)
        * (1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)))
        for base_lr in self.base_lrs
    ]