code: def test_one_iteration(world_size, test_case, fsdp_config):
      """Test FSDP with uneven divide of parameter shards."""
      if torch_version() < (1, 6, 0):
          pytest.skip("older pytorch doesn't support reduce_scatter")
  
      if world_size > torch.cuda.device_count():
          pytest.skip("Not enough GPUs.")
  
      temp_file_name = tempfile.mkstemp()[1]
      unused = tempfile.mkstemp()[1]
  
      # TODO (Min): we may want to extend this to a simple 2 layer model so that it covers
      #             more cases in FSDP. Also, assert_ref_out can be extended to multiple
      #             iterations. This could be a good bootcamp task. I should file a github
      #             issue once we merge.
      model = Linear(3, 3, bias=False)
      mp.spawn(
          _test_func,
          args=(
              world_size,
              model,
              fsdp_config,
              temp_file_name,
              unused,
              test_case,
          ),
          nprocs=world_size,
          join=True,
      )
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_uneven.py
imports: import tempfile
  import pytest
  import torch
  from torch import Tensor
  import torch.multiprocessing
  from torch.nn import Linear, Sequential
  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.fully_sharded_data_parallel import TrainingState
label_desc: Use image node_id 2 for calling a global function with example usage: test_one_iteration(world_size, test_case, fsdp_config) without return types