code: def _test_func(
      rank,
      world_size,
      tempfile_name,
      unused,
      flatten,
      mixed_precision,
      amp_context,
      half_input,
      fsdp_wrap_ckpt,
  ):
      result = dist_init(rank, world_size, tempfile_name, unused)
      assert result, "Dist init failed"
  
      # Keep initialization deterministic.
      torch.manual_seed(0)
  
      model = FSDP(
          SimpleModuleWithCheckpointing(
              flatten, mixed_precision, fsdp_wrap_ckpt
          ).cuda(),
          flatten_parameters=flatten,
          mixed_precision=mixed_precision,
      )
      optim = torch.optim.SGD(
          model.parameters(), lr=0.001, momentum=0.9
      )
  
      # Collect parameter sizes to ensure these stay consistent through the steps below.
      expected_param_shapes = {
          name: tuple(param.shape)
          for name, param in model.named_parameters()
      }
  
      # For clarity, this is what `expected_param_shapes` should look like depending on world size:
      if not flatten:
          assert expected_param_shapes == {
              "ffn.0.weight": (5,),
              "ffn.0.bias": (2,),
              "ffn.1.weight": (5,),
              "ffn.1.bias": (2,),
              "ffn.2.weight": (5,),
              "ffn.2.bias": (2,),
          }
      else:
          assert expected_param_shapes == {
              "_fsdp_wrapped_module.flat_param_0": (12,),
              "_fsdp_wrapped_module._fpw_module.ffn.1._fsdp_wrapped_module.flat_param_0": (
                  6,
              ),
          }, expected_param_shapes
  
      torch.manual_seed(1 + rank)
  
      # Train for a step.
      _train_step(
          model,
          optim,
          expected_param_shapes,
          amp_context,
          mixed_precision,
          half_input,
      )
  
      # Now do an eval step.
      _eval_step(
          model,
          optim,
          expected_param_shapes,
          amp_context,
          mixed_precision,
          half_input,
      )
  
      # And finally do another train step.
      _train_step(
          model,
          optim,
          expected_param_shapes,
          amp_context,
          mixed_precision,
          half_input,
      )
  
      teardown()
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py
imports: import contextlib
  import pytest
  import torch
  from torch import nn
  import torch.distributed
  import torch.multiprocessing
  from fairscale.fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
label_desc: Use image node_id 2 for calling a global function with example usage: _test_func(rank, world_size, tempfile_name, unused, flatten, mixed_precision, amp_context, half_input, fsdp_wrap_ckpt) without return types