code: def __init__(self):
      super().__init__()
      self.seq = nn.Sequential(
          nn.Conv2d(3, 5, 5),
          nn.AdaptiveAvgPool2d(5),
          self.Flatten(),
          nn.Linear(125, N_CLASS),
      )
  
file_path: baal/tests/ensemble_test.py
imports: import pytest
  import torch
  from torch import nn
  from torch import optim
  from torch.utils.data import Dataset
  from baal.ensemble import EnsembleModelWrapper, ensemble_prediction
label_desc: Use image node_id 1 to create a new AModel object from inherited base classes: nn with example: obj = AModel()