def test_resume_mid_epoch_warning(tmp_path):
    """Test that resuming from a mid-epoch checkpoint raises a warning unless the dataloader is stateful."""

    class NotStatefulIterable:
        def __init__(self):
            self.index = 0

        def __iter__(self):
            for i in range(self.index, len(self)):
                yield torch.ones(2, 32) * i

        def __len__(self):
            return 3

    class StatefulIterable(NotStatefulIterable):
        def state_dict(self):
            return {"index": self.index}

        def load_state_dict(self, state_dict):
            self.index = state_dict["index"]

    trainer_kwargs = {
        "default_root_dir": tmp_path,
        "accelerator": "cpu",
        "max_epochs": 1,
        "enable_model_summary": False,
        "enable_progress_bar": False,
        "logger": False,
    }

    def train_and_resume(dataloader, resume_step, expected_warning):
        # Initial training
        checkpoint_dir = tmp_path / "checkpoints"
        trainer = Trainer(
            **trainer_kwargs,
            callbacks=ModelCheckpoint(dirpath=checkpoint_dir, every_n_train_steps=1, save_top_k=-1),
        )
        trainer.fit(BoringModel(), dataloader)

        # Resume
        trainer = Trainer(**trainer_kwargs, enable_checkpointing=False)
        resume_from = checkpoint_dir / f"epoch=0-step={resume_step}.ckpt"
        warn_assert = pytest.warns if expected_warning else no_warning_call
        with warn_assert(PossibleUserWarning, match="resuming from a checkpoint that ended before"):
            trainer.fit(BoringModel(), dataloader, ckpt_path=resume_from)

    # Resume mid-epoch, no stateful dataloader -> warning
    train_and_resume(dataloader=NotStatefulIterable(), resume_step=1, expected_warning=True)

    # Resume end-of-epoch, no stateful dataloader -> no warning
    train_and_resume(dataloader=NotStatefulIterable(), resume_step=3, expected_warning=False)

    # Resume mid-epoch, stateful dataloader -> no warning
    train_and_resume(dataloader=StatefulIterable(), resume_step=1, expected_warning=False)
Test that resuming from a mid-epoch checkpoint raises a warning unless the dataloader is stateful.
tests/tests_pytorch/loops/test_training_epoch_loop.py