code: def init_model():
      model = FSDP(torch.nn.Linear(4, 4).cuda(dist.get_rank()))
      optim = torch.optim.Adam(model.parameters(), lr=0.1)
      model(torch.rand(4, 4)).sum().backward()
      optim.step()
  
      return model, optim
  
file_path: pytorch/torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py
imports: import os
  import shutil
  import torch
  import torch.distributed
  import torch.distributed.checkpoint
  import torch.multiprocessing
  from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
label_desc: Use image node_id 2 for calling a global function with example usage: init_model() and returns: model,  optim