code: def __init__(self, n_classes):
      super(LeNet5, self).__init__()
  
      self.feature_extractor = torch.nn.Sequential(
          torch.nn.Conv2d(
              in_channels=1, out_channels=6, kernel_size=5, stride=1
          ),
          torch.nn.Tanh(),
          torch.nn.AvgPool2d(kernel_size=2),
          torch.nn.Conv2d(
              in_channels=6, out_channels=16, kernel_size=5, stride=1
          ),
          torch.nn.Tanh(),
          torch.nn.AvgPool2d(kernel_size=2),
          torch.nn.Conv2d(
              in_channels=16, out_channels=120, kernel_size=5, stride=1
          ),
          torch.nn.Tanh(),
      )
  
      self.classifier = torch.nn.Sequential(
          torch.nn.Linear(in_features=120, out_features=84),
          torch.nn.Tanh(),
          torch.nn.Linear(in_features=84, out_features=n_classes),
      )
  
file_path: DeepSpeed/tests/unit/profiling/flops_profiler/test_flops_profiler.py
imports: import torch
  import pytest
  import deepspeed
  from deepspeed.profiling.flops_profiler import get_model_profile
  from unit.simple_model import SimpleModel, random_dataloader
  from unit.common import DistributedTest
  from deepspeed.runtime.utils import required_torch_version
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 1 to create a new LeNet5 object from inherited base classes: torch.nn with example: obj = LeNet5(n_classes)