code: def __init__(
      self,
      with_fsdp,
      freeze_after_wrap_fsdp,
      disable_autograd,
      fsdp_kwargs,
  ):
      super().__init__()
      self.trunk = nn.Sequential(
          self._create_block(3, 64, with_fsdp, freeze_after_wrap_fsdp),
          self._create_block(64, 64, with_fsdp, freeze_after_wrap_fsdp),
      )
      self.head = nn.Sequential(
          nn.AdaptiveAvgPool2d(output_size=(1, 1)),
          nn.Flatten(),
          nn.Linear(64, 10),
      )
      if with_fsdp and freeze_after_wrap_fsdp:
          self.fsdp_wrap(fsdp_kwargs)
      self.autograd_ctx = (
          torch.no_grad if disable_autograd else contextlib.nullcontext
      )
  
imports: import contextlib
  import sys
  from enum import Enum
  import torch
  import torch.nn
  import torch.optim
  from torch import distributed
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest, get_full_params
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 1 to create a new NestedTrunkModel object from inherited base classes: nn with example: obj = NestedTrunkModel(with_fsdp, freeze_after_wrap_fsdp, disable_autograd, fsdp_kwargs)