def test_loops_restore(tmp_path):
    """Test that required loop state_dict is loaded correctly by checkpoint connector."""
    model = BoringModel()
    checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, save_last=True)
    trainer_args = {
        "default_root_dir": tmp_path,
        "max_epochs": 1,
        "limit_train_batches": 1,
        "limit_val_batches": 1,
        "logger": False,
        "callbacks": [checkpoint_callback],
        "num_sanity_val_steps": 0,
    }
    trainer = Trainer(**trainer_args)
    trainer.fit(model)

    ckpt_path = str(tmp_path / "last.ckpt")

    trainer = Trainer(**trainer_args)
    trainer.strategy.connect(model)

    trainer_fns = list(TrainerFn)
    for fn in trainer_fns:
        trainer_fn = getattr(trainer, f"{fn.value}_loop")
        trainer_fn.load_state_dict = mock.Mock()

    for fn in trainer_fns:
        trainer.state.fn = fn
        trainer._checkpoint_connector.resume_start(ckpt_path)
        trainer._checkpoint_connector.restore_loops()

        trainer_loop = getattr(trainer, f"{fn.value}_loop")
        trainer_loop.load_state_dict.assert_called()
        trainer_loop.load_state_dict.reset_mock()

        for fn2 in trainer_fns:
            if fn2 != fn:
                trainer_loop2 = getattr(trainer, f"{fn2.value}_loop")
                trainer_loop2.load_state_dict.assert_not_called()
Test that required loop state_dict is loaded correctly by checkpoint connector.
tests/tests_pytorch/trainer/connectors/test_checkpoint_connector.py