class_docstr: Model to test FSDP(checkpoint()).
class_docstr_tok: ["Model","to","test","FSDP","(","checkpoint","(",")",")","."]
code: def forward(self, x):
      if isinstance(x, torch.Tensor):
          return self.head(self.block2(self.block1(x)))
      elif isinstance(x, list):
          ys = [self.head(self.block2(self.block1(e))) for e in x]
          return torch.cat(ys, dim=0)
  
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 2 for calling the Model obj's underlying member method code with example usage: obj.forward(x) and returns: self, torch