code: def __init__(self):
      super().__init__()
  
      # Tunable parameters
      self.hidden_size_1 = args.hidden_size_1
      self.hidden_size_2 = args.hidden_size_2
      self.learning_rate = args.learning_rate
      self.dropout = args.dropout
      self.batch_size = args.batch_size
  
      # Set class attributes
      self.data_dir = PATH_DATASETS
  
      # Hardcode some dataset specific attributes
      self.num_classes = 10
      self.dims = (1, 28, 28)
      channels, width, height = self.dims
      self.transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,)),
          ]
      )
  
      # Create a PyTorch model
      layers = [nn.Flatten()]
      width = channels * width * height
      hidden_layers = [self.hidden_size_1, self.hidden_size_2]
      num_params = 0
      for hidden_size in hidden_layers:
          if hidden_size > 0:
              layers.append(nn.Linear(width, hidden_size))
              layers.append(nn.ReLU())
              layers.append(nn.Dropout(self.dropout))
              num_params += width * hidden_size
              width = hidden_size
      layers.append(nn.Linear(width, self.num_classes))
      num_params += width * self.num_classes
  
      # Save the model and parameter counts
      self.num_params = num_params
      self.model = nn.Sequential(*layers)
  
imports: import argparse
  import logging
  import os
  import time
  import warnings
  import torch
  from pytorch_lightning import LightningModule, loggers, Trainer
  from torch import nn
  from torch.nn import functional
  from torch.utils.data import DataLoader
  from torchmetrics.functional.classification.accuracy import multiclass_accuracy
  from torchvision import transforms
  from torchvision.datasets import MNIST
label_desc: Use image node_id 1 to create a new MnistModel object from inherited base classes: LightningModule with example: obj = MnistModel()