code: def test_train_and_eval_with_checkpointing(
      flatten, mixed_precision, amp_context, half_input, fsdp_wrap_ckpt
  ):
      flatten = flatten == "flat"
      mixed_precision = mixed_precision == "fp16"
      amp_context = amp_context == "autocast"
      half_input = half_input == "halfin"
      fsdp_wrap_ckpt = fsdp_wrap_ckpt == "F->C"
  
      # Expecting an known bug in 4 out of 32 cases.
      if fsdp_wrap_ckpt and mixed_precision and not flatten:
          pytest.skip("known bug")
  
      world_size = 2
  
      with temp_files_ctx(2) as (temp_file_name, unused):
          mp.spawn(
              _test_func,
              args=(
                  world_size,
                  temp_file_name,
                  unused,
                  flatten,
                  mixed_precision,
                  amp_context,
                  half_input,
                  fsdp_wrap_ckpt,
              ),
              nprocs=world_size,
              join=True,
          )
  
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 1 for calling a global function with example usage: test_train_and_eval_with_checkpointing(flatten, mixed_precision, amp_context, half_input, fsdp_wrap_ckpt) without return types