code: def __init__(self):
      super().__init__()
      self.c_d1 = nn.Linear(in_features=28 * 28, out_features=128)
      self.c_d1_bn = nn.BatchNorm1d(128)
      self.c_d1_drop = nn.Dropout(0.3)
      self.c_d2 = nn.Linear(in_features=128, out_features=10)
      self.example_input_array = torch.rand(2, 1, 28, 28)
      self._loss = []
  
imports: import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from lightning.pytorch.core.module import LightningModule
  from torch.utils.data import DataLoader
  from tests_pytorch import _PATH_DATASETS
  from tests_pytorch.helpers.datasets import MNIST, AverageDataset, TrialMNIST
label_desc: Use image node_id 1 to create a new ParityModuleMNIST object from inherited base classes: LightningModule with example: obj = ParityModuleMNIST()