code: def test_tp_checkpoint_load_on_meta_device(self):
      CHECKPOINT_DIR = self.temp_dir
      mesh_shpe = (self.world_size,)
      tp_mesh = init_device_mesh(self.device_type, mesh_shpe)
  
      # create model and move it to GPU with id rank
      model = UnevenShardedModel(self.device_type).cuda(self.rank)
      # Parallelize the module based on the given Parallel Style.
      parallelize_plan = {
          "net1": RowwiseParallel(),
          "net2": ColwiseParallel(),
      }
      model = parallelize_module(
          model, tp_mesh, parallelize_plan=parallelize_plan
      )
      original_state_dict = deepcopy(model.state_dict())
  
      DCP.save_state_dict(
          state_dict=original_state_dict,
          storage_writer=DCP.FileSystemWriter(CHECKPOINT_DIR),
      )
  
      model2 = parallelize_module(
          UnevenShardedModel("meta"),
          tp_mesh,
          parallelize_plan=parallelize_plan,
      )
      state_dict_to_load = model2.state_dict()
  
      DCP.load_state_dict(
          state_dict=state_dict_to_load,
          storage_reader=DCP.FileSystemReader(CHECKPOINT_DIR),
      )
      model2.load_state_dict(state_dict_to_load, assign=True)
      state_dict_after_load = model2.state_dict()
  
      # After loading, check whether params in state_dict_after_load are equal to original_state_dict.
      for param1, param2 in zip(
          original_state_dict.values(), state_dict_after_load.values()
      ):
          self.assertEqual(param1, param2)
  
file_path: pytorch/test/distributed/checkpoint/test_tp_checkpoint.py
imports: from copy import deepcopy
  import torch
  import torch.distributed.checkpoint
  from torch.distributed._tensor import init_device_mesh
  from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, RowwiseParallel
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, MLPModule, skip_if_lt_x_gpu, with_comms
  from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
label_desc: Use image node_id 2 for calling the TestTpCheckpoint obj's underlying member method code with example usage: obj.test_tp_checkpoint_load_on_meta_device() without return types