def test_lazy_load_full_state_dict(_, lazy_load_mock, torch_load_mock, tmp_path):
    """Test that loading a single file (full state) is lazy to reduce peak CPU memory usage."""
    model = BoringModel()
    checkpoint = {"state_dict": model.state_dict()}
    lazy_load_mock.return_value = checkpoint

    strategy = FSDPStrategy()
    trainer = Trainer()
    model.trainer = trainer
    strategy._lightning_module = model
    strategy.model = model

    file = tmp_path / "test.ckpt"
    file.touch()

    strategy.load_checkpoint(checkpoint_path=file)
    lazy_load_mock.assert_called_once()
Test that loading a single file (full state) is lazy to reduce peak CPU memory usage.
tests/tests_pytorch/strategies/test_fsdp.py