code: def test_get_full_state_dict_context_offload(
      set_type_mock, monkeypatch, torch_ge_2_1
  ):
      """Test that the state dict context manager handles CPU offloading depending on the PyTorch version."""
      monkeypatch.setattr(
          "lightning.fabric.strategies.fsdp._TORCH_GREATER_EQUAL_2_1",
          torch_ge_2_1,
      )
  
      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
              is torch_ge_2_1
          )  # model config
          assert (
              set_type_mock.call_args_list[0][0][3].offload_to_cpu
              is torch_ge_2_1
          )  # 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
  
imports: from datetime import timedelta
  from re import escape
  from unittest import mock
  from unittest.mock import ANY, MagicMock, Mock
  import lightning.fabric
  import pytest
  import torch
  import torch.nn
  from lightning.fabric.plugins import HalfPrecision
  from lightning.fabric.plugins.environments import LightningEnvironment
  from lightning.fabric.strategies import FSDPStrategy
  from lightning.fabric.strategies.fsdp import _FSDPBackwardSyncControl, _get_full_state_dict_context, _has_meta_device_parameters, _is_sharded_checkpoint
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
  from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, FullyShardedDataParallel, MixedPrecision
  from torch.optim import Adam
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 24 for calling a global function with example usage: test_get_full_state_dict_context_offload(set_type_mock, monkeypatch, torch_ge_2_1) without return types