def METHOD_NAME(self, training_phase: TrainingPhase = TrainingPhase.BACKWARD):
    old_training_phase = self._training_phase
    try:
        self._training_phase = training_phase
        yield
    finally:
        self._training_phase = old_training_phase