code: def test_input_type(temp_files, fsdp_config, input_cls):
      """Test FSDP with input being a list or a dict, only single GPU."""
  
      if torch_version() < (1, 7, 0):
          # This test runs multiple test cases in a single process. On 1.6.0 it
          # throw an error like this:
          #     RuntimeError: Container is already initialized! Cannot initialize it twice!
          pytest.skip(
              "older pytorch doesn't work well with single process dist_init multiple times"
          )
  
      result = dist_init(
          rank=0,
          world_size=1,
          filename=temp_files[0],
          filename_rpc=temp_files[1],
      )
      assert result, "Dist init failed"
  
      assert isinstance(fsdp_config, dict), str(fsdp_config)
  
      class Model(Module):
          def __init__(self):
              super().__init__()
              self.layer = Linear(4, 4)
  
          def forward(self, input):
              if isinstance(input, list):
                  input = input[0]
              else:
                  assert isinstance(input, dict), input
                  input = input["in"]
              return self.layer(input)
  
      model = FSDP(Model(), **fsdp_config).cuda()
      optim = SGD(model.parameters(), lr=0.1)
  
      for _ in range(5):
          in_data = torch.rand(64, 4).cuda()
          in_data.requires_grad = True
          if input_cls is list:
              in_data = [in_data]
          else:
              assert input_cls is dict
              in_data = {"in": in_data}
  
          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_input.py
imports: import tempfile
  import pytest
  import torch
  from torch.nn import Linear, Module
  from torch.optim import SGD
  from fairscale.fair_dev.testing.testing import dist_init, rmf, 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 2 for calling a global function with example usage: test_input_type(temp_files, fsdp_config, input_cls) and returns: self