code: def __init__(self, feature_extraction=False) -> None:
      super().__init__()
      self.feature_extraction = feature_extraction
  
      self.conv1 = nn.Conv2d(3, 3, 5)
      self.relu1 = nn.ReLU()
      self.pool1 = nn.MaxPool2d(2, 2)
  
      if not self.feature_extraction:
          self.conv2 = nn.Conv2d(3, 10, 2)
  
imports: import unittest
  from typing import Callable, List, Union
  import torch
  import torch.nn
  from captum.insights import AttributionVisualizer, Batch
  from captum.insights.attr_vis.app import FilterConfig
  from captum.insights.attr_vis.features import BaseFeature, FeatureOutput, ImageFeature
  from tests.helpers.basic import BaseTest
label_desc: Use image node_id 1 to create a new TinyCnn object from inherited base classes: nn with example: obj = TinyCnn(feature_extraction)