def test_get_full_state_dict_context_offload(set_type_mock, monkeypatch):
    """Test that the state dict context manager handles CPU offloading."""

    with _get_full_state_dict_context(module=Mock(spec=FullyShardedDataParallel), world_size=1):
        assert set_type_mock.call_args_list[0][0][2].offload_to_cpu  # model config
        assert set_type_mock.call_args_list[0][0][3].offload_to_cpu  # optim config

    set_type_mock.reset_mock()

    with _get_full_state_dict_context(module=Mock(spec=FullyShardedDataParallel), world_size=4):
        assert set_type_mock.call_args_list[0][0][2].offload_to_cpu  # model config
        assert set_type_mock.call_args_list[0][0][3].offload_to_cpu  # optim config
Test that the state dict context manager handles CPU offloading.
tests/tests_fabric/strategies/test_fsdp.py