def test_restarting_mid_epoch_raises_warning(tmp_path, stop_in_the_middle, model_cls):
    """Test that a warning is raised if training is restarted from mid-epoch."""
    limit_train_batches = 8
    trainer_kwargs = {
        "default_root_dir": tmp_path,
        "limit_train_batches": limit_train_batches,
        "limit_val_batches": 0,
        "enable_progress_bar": False,
        "enable_model_summary": False,
    }
    trainer = Trainer(max_epochs=1, **trainer_kwargs)
    model = model_cls(limit_train_batches // 2 if stop_in_the_middle else -1)

    if stop_in_the_middle:
        with pytest.raises(CustomException):
            trainer.fit(model)
    else:
        trainer.fit(model)

    ckpt_path = str(tmp_path / "resume.ckpt")
    trainer.save_checkpoint(ckpt_path)

    trainer = Trainer(max_epochs=2, **trainer_kwargs)
    model.stop_batch_idx = -1

    context_manager = pytest.warns if stop_in_the_middle else no_warning_call
    with context_manager(UserWarning, match="resuming from a checkpoint that ended"):
        trainer.fit(model, ckpt_path=ckpt_path)

    if stop_in_the_middle:
        trainer = Trainer(max_epochs=2, **trainer_kwargs)
        with pytest.warns(UserWarning, match="resuming from a checkpoint that ended"):
            trainer.fit(model, ckpt_path=ckpt_path)
Test that a warning is raised if training is restarted from mid-epoch.
tests/tests_pytorch/models/test_restore.py