code: def test_single_run():
      if sys.version_info.major == 3 and sys.version_info.minor > 10:
          pytest.skip("torch.fx doesn't seem to work 3.11 yet")
      if torch_version() < (1, 8, 0):
          pytest.skip("requires torch version >= 1.8.0")
      from fairscale.experimental.nn.auto_shard import shard_model
  
      model = TransformerModel(
          ntokens, emsize, nhead, nhid, nlayers, dropout
      )
      sharded_model = shard_model(model)
      assert (
          len(sharded_model) == 2
      ), "Length is sharded model is incorrect."
      expected_param_nums = [5998600, 5785383]
      for i, model in enumerate(sharded_model):
          param_count = {}
          for name, module in model.named_modules():
              if "." in name:
                  continue
  
              param_count[name] = sum(
                  [x.numel() for x in module.parameters()]
              )
          assert expected_param_nums[i] == param_count[""]
  
      src_mask = torch.randn((35, 35), dtype=torch.float32)
      src = torch.randint(1, ntokens, (35, 20))
      input = [src, src_mask]
      for model in sharded_model:
          if type(input) == list:
              input = model(*input)
          else:
              input = model(input)
  
      assert input.size() == torch.Size([35, 20, 28783])
  
imports: import math
  import sys
  import pytest
  import torch
  import torch.nn
  import torch.nn
  from fairscale.internal import torch_version
label_desc: Use image node_id 1 for calling a global function with example usage: test_single_run() without return types