code: def test_dtensor_sharded_model_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,
          state_dict_config=ShardedStateDictConfig(
              offload_to_cpu=offload_to_cpu
          ),
      )
  
      checkpoint = io.BytesIO()
      torch.save(model.state_dict(), checkpoint)
      # Deepcopy to save current state_dict to compare with the state_dict loaded back below.
      ref_state_dict = deepcopy(model.state_dict())
  
      # Update the parameters so model.state_dict() will be different from ref_dtensor_sd.
      model(model.get_input()).sum().backward()
      optim.step()
  
      # Load ref_state_dict back.
      checkpoint.seek(0)
      load_ref_state_dict = torch.load(checkpoint)
      model.load_state_dict(load_ref_state_dict)
      new_state_dict = model.state_dict()
  
      # Check whether new_state_dict is the same as ref_state_dict.
      for (k1, v1), (k2, v2) in zip(
          ref_state_dict.items(), new_state_dict.items()
      ):
          # check whether fqn are the same
          self.assertEqual(k1, k2)
  
          self.assertEqual(type(v1), DTensor)
          self.assertEqual(type(v2), DTensor)
          # check whether DTensor are the same
          self.assertEqual(v1, v2)
  
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 5 for calling the TestFSDPWithDeviceMeshAndDTensor obj's underlying member method code with example usage: obj.test_dtensor_sharded_model_load_state_dict(offload_to_cpu, is_even_sharded_model) without return types