code: def test_fsdp_init_with_device_mesh(self, 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 = model.state_dict()
      optim_state_dict = FSDP.optim_state_dict(model, optim)
  
      for v in state_dict.values():
          self.assertEqual(type(v), DTensor)
          self.assertEqual(len(v.placements), 1)
          self.assertEqual(v.placements[0], (Shard(dim=0)))
          self.assertEqual(v.device_mesh, device_mesh)
  
      for state in optim_state_dict["state"].values():
          for k, v in state.items():
              if k != "step":
                  self.assertEqual(type(v), DTensor)
                  self.assertEqual(len(v.placements), 1)
                  self.assertEqual(v.placements[0], (Shard(dim=0)))
                  self.assertEqual(v.device_mesh, device_mesh)
  
      state_dict_type = FSDP.get_state_dict_type(model)
      # If device_mesh is used when initializing FSDP, the field _use_dtensor will
      # automatically be set to True if StateDictType is set to SHARDED_STATE_DICT.
      self.assertEqual(
          state_dict_type.state_dict_config._use_dtensor, True
      )
      self.assertEqual(
          state_dict_type.optim_state_dict_config._use_dtensor, True
      )
  
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 2 for calling the TestFSDPWithDeviceMeshAndDTensor obj's underlying member method code with example usage: obj.test_fsdp_init_with_device_mesh(is_even_sharded_model) without return types