code: def _init_model(rank, world_size):
      device_mesh = init_device_mesh(DEVICE, (world_size,))
  
      # Create a dummy model and wrap it in FSDP
      model = Model().cuda()
      device_mesh = init_device_mesh(DEVICE, (world_size,))
      model = FSDP(model, device_mesh=device_mesh, use_orig_params=True)
  
      optim = torch.optim.Adam(model.parameters(), lr=0.0001)
  
      _patch_model_state_dict(model)
      _patch_optimizer_state_dict(model, optimizers=optim)
  
      return model, optim
  
imports: import os
  import shutil
  import traceback
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  from torch.distributed._tensor.device_mesh import init_device_mesh
  from torch.distributed.checkpoint.filesystem import _FileSystemCheckpointer
  from torch.distributed.checkpoint.state_dict import _patch_model_state_dict, _patch_optimizer_state_dict
  from torch.distributed.fsdp import FullyShardedDataParallel
label_desc: Use image node_id 1 for calling a global function with example usage: _init_model(rank, world_size) and returns: model,  optim