class_docstr: Model to test FSDP(checkpoint(), checkpoint()).
class_docstr_tok: ["Model","to","test","FSDP","(","checkpoint","(",")",",","checkpoint","(",")",")","."]
code: def __init__(self):
      super().__init__()
      self.block1 = nn.Sequential(
          nn.Conv2d(3, 4, kernel_size=3),
          nn.BatchNorm2d(4),
          nn.ReLU(inplace=True),
      )
      self.block2 = nn.Sequential(
          nn.Conv2d(4, 4, kernel_size=3),
          nn.BatchNorm2d(4),
          nn.ReLU(inplace=False),
      )
      self.block3 = nn.Sequential(
          nn.Conv2d(4, 8, kernel_size=3),
          nn.BatchNorm2d(8),
          nn.ReLU(inplace=True),
      )
      self.head = nn.Sequential(
          nn.AdaptiveAvgPool2d(output_size=(1, 1)),
          nn.Flatten(),
          nn.Linear(8, 10),
      )
  
imports: import contextlib
  import pickle
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  import torch.optim
  from fairscale.fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.internal import torch_version
  from fairscale.nn import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import auto_wrap_bn
  from fairscale.nn.wrap import enable_wrap, wrap
label_desc: Use image node_id 1 to create a new Model2 object from inherited base classes: nn with example: obj = Model2()