def test_should_stop_mid_epoch(tmp_path):
    """Test that training correctly stops mid epoch and that validation is still called at the right time."""

    class TestModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.validation_called_at = None

        def training_step(self, batch, batch_idx):
            if batch_idx == 4:
                self.trainer.should_stop = True
            return super().training_step(batch, batch_idx)

        def validation_step(self, *args):
            self.validation_called_at = (self.trainer.current_epoch, self.trainer.global_step)
            return super().validation_step(*args)

    model = TestModel()
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=1, limit_train_batches=10, limit_val_batches=1)
    trainer.fit(model)

    # even though we stopped mid epoch, the fit loop finished normally and the current epoch was increased
    assert trainer.current_epoch == 1
    assert trainer.global_step == 5
    assert model.validation_called_at == (0, 5)
Test that training correctly stops mid epoch and that validation is still called at the right time.
tests/tests_pytorch/loops/test_training_loop.py