code: def __init__(
      self, embedding_size: int, with_fsdp: bool, process_group
  ):
      super().__init__()
      self.conv1 = self._conv_block(3, embedding_size)
      self.conv2: nn.Module = self._conv_block(
          embedding_size, embedding_size // 2
      )
      self.conv3: nn.Module = self._conv_block(
          embedding_size // 2, embedding_size
      )
      self.pool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
      self.flatten = nn.Flatten(start_dim=1)
      self.relu = nn.ReLU()
      self.fc1: nn.Module = nn.Linear(
          embedding_size, 2 * embedding_size
      )
      self.fc2: nn.Module = nn.Linear(
          2 * embedding_size, 2 * embedding_size
      )
      self.fc3: nn.Module = nn.Linear(
          2 * embedding_size, embedding_size + 1
      )
      self.fc4: nn.Module = nn.Linear(
          embedding_size + 1, embedding_size
      )
      if with_fsdp:
          self.conv2 = FullyShardedDataParallel(
              self.conv2, process_group=process_group
          )
          self.conv3 = FullyShardedDataParallel(
              self.conv3,
              process_group=process_group,
              flatten_parameters=False,
          )
          self.fc1 = FullyShardedDataParallel(
              self.fc1, process_group=process_group
          )
          self.fc3 = FullyShardedDataParallel(
              self.fc3,
              process_group=process_group,
              flatten_parameters=False,
          )
  
imports: import functools
  import os
  import tempfile
  from parameterized import parameterized
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.optim import Adam
  from fairscale.fair_dev.testing.testing import in_temporary_directory, skip_if_single_gpu, temp_files_ctx
  from fairscale.nn import FullyShardedDataParallel
  from tests.nn.data_parallel.test_fsdp import DistributedTest, MixtureOfExperts, rename_test, spawn_and_init
label_desc: Use image node_id 1 to create a new ConvolutionalModel object from inherited base classes: nn with example: obj = ConvolutionalModel(embedding_size, with_fsdp, process_group)