code: def _dist_worker(
      rank, world_size, files, outer_flat, inner_flat, sharing
  ):
      # Get data from files.
      file1, file2, sd_before, sd_after, in_data = files
      sd_before = torch.load(
          sd_before,
          map_location=lambda storage, loc: storage.cuda(rank),
      )
      sd_after = torch.load(
          sd_after, map_location=lambda storage, loc: storage.cuda(rank)
      )
      in_data = torch.load(
          in_data, map_location=lambda storage, loc: storage.cuda(rank)
      )
  
      result = dist_init(
          rank=rank,
          world_size=world_size,
          filename=file1,
          filename_rpc=file2,
      )
      assert result, "Dist init failed"
  
      fsdp_model = FSDP(
          Model(with_fsdp=True, inner_flat=inner_flat, sharing=sharing),
          flatten_parameters=outer_flat,
      )
      fsdp_model.load_state_dict(sd_before)
  
      _train(fsdp_model, in_data)
  
      objects_are_equal(
          sd_after, fsdp_model.state_dict(), raise_exception=True
      )
  
      teardown()
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_shared_weights.py
imports: from copy import deepcopy
  import pytest
  import torch
  import torch.multiprocessing
  from torch.nn import Linear, Module
  from torch.optim import SGD
  from fairscale.fair_dev.testing.testing import dist_init, objects_are_equal, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.nn.data_parallel import FullyShardedDataParallel
label_desc: Use image node_id 3 for calling a global function with example usage: _dist_worker(rank, world_size, files, outer_flat, inner_flat, sharing) without return types