code: def _distributed_worker(
      gpu_id,
      world_size,
      with_fsdp,
      with_checkpoint,
      filename,
      filename_rpc,
      expected,
      model_hidden_dim,
      fsdp_config,
  ):
      torch.cuda.set_device(gpu_id)
  
      rank = gpu_id
      result = dist_init(rank, world_size, filename, filename_rpc)
      assert result, "Dist init failed"
  
      torch.manual_seed(0)
      torch.backends.cudnn.deterministic = True
  
      # Note that FSDP auto-cast the input in AMP mode. So we don't need to call half() here.
      batch = torch.randn(size=(2, 3, 224, 224)).cuda()
  
      model = create_model(
          with_fsdp, with_checkpoint, model_hidden_dim, fsdp_config
      )
      model = model.cuda()
      if with_fsdp:
          model = to_fsdp(model, fsdp_config)
      else:
          model = DistributedDataParallel(
              model, device_ids=[gpu_id], bucket_cap_mb=500
          )
  
      # We enable momentum so that after the first iteration, the optimizer state is added
      # to the total memory used.
      criterion = nn.MSELoss()
      optimizer = optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
  
      # Set AMP context if needed.
      context = contextlib.suppress()
      if (
          "mixed_precision" in fsdp_config
          and fsdp_config["mixed_precision"]
      ):
          context = torch.cuda.amp.autocast(enabled=True)
  
      # We have observed that sometimes after 3rd iteration, 4th one can fail (not on this
      # test but on much bigger scale tests). We run 4 iterations here just in case it happens.
      iterations = 4
  
      results = {}  # results of memory stats
      for iteration in range(iterations):
          get_cur_mem(gpu_id, results, f"iter {iteration}: start")
  
          with context:
              out = model(batch)
              get_cur_mem(
                  gpu_id, results, f"iter {iteration}: after fwd"
              )
  
              out = sum(o.sum() for o in out[0])
              fake_loss = criterion(out, torch.tensor(0.0).cuda())
              get_cur_mem(
                  gpu_id, results, f"iter {iteration}: after loss"
              )
  
          fake_loss.backward()
          get_cur_mem(gpu_id, results, f"iter {iteration}: after bwd")
  
          optimizer.step()
          get_cur_mem(gpu_id, results, f"iter {iteration}: after step")
  
          # It is important to use `set_to_none` below, not optimizer.zero_grad() to reclaim memory.
          if torch_version() >= (1, 7, 0):
              model.zero_grad(set_to_none=True)
          else:
              for p in model.parameters():
                  p.grad = None
          get_cur_mem(gpu_id, results, f"iter {iteration}: done")
  
      dump_all_tensors(gpu_id)
      print(results)
  
      def cmp(results, expected):
          ret = ""
          assert (
              results.keys() == expected.keys()
          ), f"{list(results.keys())} vs. {list(expected.keys())}"
          for k, v in results.items():
              exp = expected[k]
              if abs(exp - v) > 1:  # allow 1MB rounding differences
                  ret += f"{k}: got {v}, expected {exp}\n"
          return ret
  
      output = cmp(results, expected)
      assert not output, output
  
      teardown()
  
imports: import contextlib
  import pytest
  import torch
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  import torch.optim
  from fairscale.fair_dev.testing.testing import dist_init, dump_all_tensors, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.internal import torch_version
  from fairscale.internal.parallel import get_process_group_cached
  from fairscale.nn import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import auto_wrap_bn
label_desc: Use image node_id 4 for calling a global function with example usage: _distributed_worker(gpu_id, world_size, with_fsdp, with_checkpoint, filename, filename_rpc, expected, model_hidden_dim, fsdp_config) and returns: ret