code: def __init__(self, with_fsdp=False, inner_flat=False, sharing=None):
      super().__init__()
      self.l0 = Linear(4, 4, bias=True).cuda()
      self.l1 = Linear(4, 4, bias=True).cuda()
      self.l2 = Linear(4, 4, bias=True).cuda()
      self.l3 = Linear(4, 4, bias=True).cuda()
  
      # share the weights. the layer must have at least 1 param is that's not
      # shared. Therefore, we have bias=True and testing either sharing the
      # weight or the bias.
      if sharing == "share_only_weights":
          self.l1.weight = self.l3.weight
      elif sharing == "share_only_bias":
          self.l1.bias = self.l3.bias
      else:
          assert sharing is None or sharing == "share_none"
  
      if with_fsdp:
          # Shared layers much be un-flatten.
          self.l1 = FSDP(self.l1, flatten_parameters=False)
          self.l2 = FSDP(self.l2, flatten_parameters=inner_flat)
          self.l3 = FSDP(self.l3, flatten_parameters=False)
  
          if sharing in ["share_only_weights"]:
              self.l3.append_shared_param(self.l1.module.weight)
          if sharing in ["share_only_bias"]:
              self.l3.append_shared_param(self.l1.module.bias)
  
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 1 to create a new Model object from inherited base classes: Module with example: obj = Model(with_fsdp, inner_flat, sharing)