code: def test_2d_load_state_dict(self, is_even_sharded_model):
      simple_model = (
          SimpleModel if is_even_sharded_model else SimpleModelUneven
      )
  
      torch.manual_seed(0)
      mesh_2d = init_device_mesh(
          self.device_type,
          (2, self.world_size // 2),
          mesh_dim_names=("dp", "tp"),
      )
      tp_mesh = mesh_2d["tp"]
      dp_mesh = mesh_2d["dp"]
      parallelize_plan = {
          "net1": ColwiseParallel(),
          "net2": RowwiseParallel(),
      }
      model_2d = parallelize_module(
          simple_model().cuda(), tp_mesh, parallelize_plan
      )
      model_2d = FSDP(
          model_2d, device_mesh=dp_mesh, use_orig_params=True
      )
      optim_2d = torch.optim.Adam(model_2d.parameters(), lr=0.01)
  
      FSDP.set_state_dict_type(
          model_2d,
          StateDictType.SHARDED_STATE_DICT,
      )
      checkpoint = io.BytesIO()
      torch.save(model_2d.state_dict(), checkpoint)
      # Deepcopy to save current state_dict to compare with the state_dict loaded back below.
      ref_state_dict = deepcopy(model_2d.state_dict())
  
      # Update the parameters so model.state_dict() will be different from ref_dtensor_sd.
      model_2d(model_2d.get_input().cuda(self.rank)).sum().backward()
      optim_2d.step()
  
      # Load ref_state_dict back.
      checkpoint.seek(0)
      load_ref_state_dict = torch.load(checkpoint)
      model_2d.load_state_dict(load_ref_state_dict)
      new_state_dict = model_2d.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), DT)
          self.assertEqual(type(v2), DT)
          # check whether DTensor are the same
          # TODO: 2D DTensor comparison is not supported at the time, so we are comparing the spec and the local tensor for now.
          # TODO: Update it to compare the two DTensors once 2D DTensor comparison is supported.
          self.assertEqual(v1.to_local(), v2.to_local())
          self.assertEqual(v1.device_mesh, v2.device_mesh)
          self.assertEqual(v1.placements, v2.placements)
  
imports: import io
  from copy import deepcopy
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  from torch.distributed._tensor import DTensor, init_device_mesh, Replicate, Shard
  from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict, set_optimizer_state_dict
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp._common_utils import _get_module_fsdp_state, clean_tensor_name
  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, RowwiseParallel
  from torch.distributed.tensor.parallel.fsdp import DTensorExtensions
  from torch.distributed.tensor.parallel.input_reshard import input_reshard
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms
label_desc: Use image node_id 3 for calling the TestNew2dParallelStateDict obj's underlying member method code with example usage: obj.test_2d_load_state_dict(is_even_sharded_model) without return types