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 Model(Module):
          def __init__(self):
              super().__init__()
              self.inner = FSDP(Linear(4, 4), **fsdp_config)
              self.outer = Linear(4, 5)
  
          def forward(self, x):
              # Forward twice.
              i = self.inner(x)
              j = self.inner(x)
              return self.outer(i + j)
  
      model = FSDP(Model(), **fsdp_config).cuda()
      optim = SGD(model.parameters(), lr=0.1)
  
      for _ in range(3):
          in_data = torch.rand(64, 4).cuda()
          in_data.requires_grad = True
          out = model(in_data)
          out.sum().backward()
          optim.step()
          optim.zero_grad()
  
      model.assert_state(TrainingState.IDLE)
      teardown()
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_multiple_forward.py
imports: import tempfile
  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, skip_if_single_gpu, 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