def METHOD_NAME(self, trainer):
    if trainer.mode in [ModeKeys.TRAIN, ModeKeys.EVAL]:
        epoch = trainer.epoch + 1
    else:
        raise ValueError(
            f'trainer mode should be {ModeKeys.TRAIN} or {ModeKeys.EVAL}, '
            f'but got {trainer.mode}')
    return epoch