code: def test_dtensor_sharded_optim_load_state_dict(
      self, offload_to_cpu, is_even_sharded_model
  ):
      device_mesh = init_device_mesh(
          self.device_type, (self.world_size,)
      )
      model, optim = self._create_model(
          is_even_sharded_model, device_mesh
      )
  
      FSDP.set_state_dict_type(
          model,
          StateDictType.SHARDED_STATE_DICT,
          optim_state_dict_config=ShardedOptimStateDictConfig(
              offload_to_cpu=offload_to_cpu
          ),
      )
  
      checkpoint = io.BytesIO()
      torch.save(FSDP.optim_state_dict(model, optim), checkpoint)
      # Deepcopy to save current optim_state_dict to compare with the optim_state_dict loaded back below.
      ref_optim_state_dict = deepcopy(
          FSDP.optim_state_dict(model, optim)
      )
  
      # Update the parameters so FSDP.optim_state_dict() will be different from ref_optim_state_dict.
      model(model.get_input()).sum().backward()
      optim.step()
  
      # Load ref_optim_state_dict back.
      checkpoint.seek(0)
      load_ref_optim_state_dict = torch.load(checkpoint)
      optim.load_state_dict(
          FSDP.optim_state_dict_to_load(
              model, optim, load_ref_optim_state_dict
          )
      )
      new_optim_state_dict = FSDP.optim_state_dict(model, optim)
  
      # Check whether new_optim_state_dict is the same as ref_optim_state_dict.
      for new_optim_state_dict_item, ref_optim_state_dict_item in zip(
          new_optim_state_dict["state"].items(),
          ref_optim_state_dict["state"].items(),
      ):
          # check FQN are the same
          self.assertEqual(
              new_optim_state_dict_item[0], ref_optim_state_dict_item[0]
          )
          for new_optim_hyper_param, ref_optim_hyper_param in zip(
              new_optim_state_dict_item[1].items(),
              ref_optim_state_dict_item[1].items(),
          ):
              k1, v1 = new_optim_hyper_param
              k2, v2 = ref_optim_hyper_param
  
              # check whether keys are the same
              self.assertEqual(k1, k2)
              # check whether values are the same
              self.assertEqual(v1, v2)
  
              if k1 != "step":
                  self.assertEqual(type(v1), DTensor)
                  self.assertEqual(type(v2), DTensor)
  
imports: import io
  from copy import deepcopy
  import torch
  import torch.nn
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed._tensor import DTensor, Shard
  from torch.distributed.device_mesh import init_device_mesh
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.api import ShardedOptimStateDictConfig, ShardedStateDictConfig, StateDictType
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, skip_if_lt_x_gpu, with_comms
label_desc: Use image node_id 4 for calling the TestFSDPWithDeviceMeshAndDTensor obj's underlying member method code with example usage: obj.test_dtensor_sharded_optim_load_state_dict(offload_to_cpu, is_even_sharded_model) without return types