code: def __init__(self, flatten, mixed_precision, fsdp_wrap_ckpt):
      super().__init__()
      if fsdp_wrap_ckpt:
          middle_module = FSDP(
              checkpoint_wrapper(nn.Linear(3, 3)),
              flatten_parameters=flatten,
              mixed_precision=mixed_precision,
          )
      else:
          middle_module = checkpoint_wrapper(
              FSDP(
                  nn.Linear(3, 3),
                  flatten_parameters=flatten,
                  mixed_precision=mixed_precision,
              )
          )
  
      self.ffn = nn.Sequential(
          nn.Linear(3, 3), middle_module, nn.Linear(3, 3)
      )
  
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 to create a new SimpleModuleWithCheckpointing object from inherited base classes: nn with example: obj = SimpleModuleWithCheckpointing(flatten, mixed_precision, fsdp_wrap_ckpt)