def on_train_epoch_end(self, trainer, model):
        """
        Check to see if new stage is reached for the next epoch, and if so, prepare the new stage by
        changing the dataloader.

        (We do next epoch so that the dataloader is prepared before the next epoch)
        """
        next_epoch = trainer.current_epoch + 1

        # Check if stage should be increased
        if next_epoch >= self.stage_epochs_cume[self._current_stage] and self._current_stage < len(self.stage_params) - 1:
            self._current_stage += 1
            self._update_to_current_stage(trainer, model)

        return super().on_train_epoch_end(trainer, model)

        Check to see if new stage is reached for the next epoch, and if so, prepare the new stage by
        changing the dataloader.

        (We do next epoch so that the dataloader is prepared before the next epoch)
        
analyze/convnexts4nd/src/callbacks/progressive_resizing.py