def METHOD_NAME(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
    if not self._skip_next_epoch_step:
        self.compression_ctrl.scheduler.epoch_step()
    self._skip_next_epoch_step = False