code: def test_shared_weight(temp_files, outer_flat, inner_flat, sharing):
      """Test FSDP with a model with shared weights."""
  
      outer_flat = outer_flat == "outer_flat"
      inner_flat = inner_flat == "inner_flat"
      world_size = 2
  
      # Get reference results.
      model = Model(sharing=sharing)
      sd_before = deepcopy(model.state_dict())
      in_data = torch.rand(1, 4).cuda()
      _train(model, in_data, world_size)
      sd_after = deepcopy(model.state_dict())
      # Before and after state should not be equal.
      assert not objects_are_equal(sd_before, sd_after)
  
      # Save data
      torch.save(sd_before, temp_files[2])
      torch.save(sd_after, temp_files[3])
      torch.save(in_data, temp_files[4])
  
      # Run FSDP
      mp.spawn(
          _dist_worker,
          (world_size, temp_files, outer_flat, inner_flat, sharing),
          nprocs=world_size,
      )
  
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 2 for calling a global function with example usage: test_shared_weight(temp_files, outer_flat, inner_flat, sharing) without return types