code: def _test_func(rank, world_size, fsdp_config, tempfile_name, unused):
      result = dist_init(rank, world_size, tempfile_name, unused)
      assert result, "Dist init failed"
  
      assert isinstance(fsdp_config, dict), str(fsdp_config)
  
      class InnerModel(Module):
          def __init__(self):
              super().__init__()
              self.layers = Sequential(
                  FSDP(Linear(5, 5), **fsdp_config),
              )
  
          def forward(self, x):
              return self.layers(x)
  
      inner_model = InnerModel()
      model = FSDP(inner_model, **fsdp_config).cuda()
      optim = SGD(model.parameters(), lr=0.1)
  
      for i in range(3):
          input = torch.rand((1, 5), dtype=torch.float).cuda()
          input.requires_grad = True
          output = model(input)
          output.sum().backward()
          optim.step()
          optim.zero_grad()
      input = torch.rand((1, 5), dtype=torch.float).cuda()
      output = model(input)
  
      model.assert_state(TrainingState.IDLE)
  
      # second time to rewrap the inner model
      rewrapped_model = FSDP(inner_model, **fsdp_config).cuda()
      rewrapped_output = rewrapped_model(input)
  
      assert torch.allclose(output, rewrapped_output)
      teardown()
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_multiple_wrapping.py
imports: import tempfile
  import pytest
  import torch
  import torch.multiprocessing
  from torch.nn import Linear, Module, Sequential
  from torch.optim import SGD
  from fairscale.fair_dev.testing.testing import dist_init, skip_if_no_cuda, teardown
  from fairscale.internal import torch_version
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import TrainingState
label_desc: Use image node_id 1 for calling a global function with example usage: _test_func(rank, world_size, fsdp_config, tempfile_name, unused) and returns: self