def test_deepspeed_load_checkpoint_optimzer_state_requested(_, optimzer_state_requested, tmp_path):
    """Test that the DeepSpeed strategy loads the optimizer state only when requested."""
    from deepspeed import DeepSpeedEngine

    strategy = DeepSpeedStrategy()
    optimizer = Mock(spec=Optimizer)
    model = Mock(spec=DeepSpeedEngine, optimizer=optimizer)
    model.modules.return_value = [model]

    # required, otherwise mock cannot be unpacked
    model.load_checkpoint.return_value = [None, {}]

    state = {"model": model}
    if optimzer_state_requested:
        state["optimizer"] = optimizer

    strategy.load_checkpoint(path=tmp_path, state=state)
    model.load_checkpoint.assert_called_with(
        tmp_path,
        tag="checkpoint",
        load_optimizer_states=optimzer_state_requested,
        load_lr_scheduler_states=False,
        load_module_strict=True,
    )
Test that the DeepSpeed strategy loads the optimizer state only when requested.
tests/tests_fabric/strategies/test_deepspeed.py