def METHOD_NAME(self, iter_: int, stage_id: int) -> None:
    lr = self._lr * _get_multiplier(self._stages[stage_id], iter_)
    for param_group in self.param_groups:
        # pyre-ignore [16]
        param_group[self._lr_param] = lr