code: def __init__(self, hidden_size=64, learning_rate=2e-4):
      super().__init__()
  
      # Set our init args as class attributes
      self.hidden_size = hidden_size
      self.learning_rate = learning_rate
  
      # 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,)),
          ]
      )
  
      # Define PyTorch model
      self.model = nn.Sequential(
          nn.Flatten(),
          nn.Linear(channels * width * height, hidden_size),
          nn.ReLU(),
          nn.Dropout(0.1),
          nn.Linear(hidden_size, hidden_size),
          nn.ReLU(),
          nn.Dropout(0.1),
          nn.Linear(hidden_size, self.num_classes),
      )
  
imports: import pytorch_lightning
  import torch
  from pytorch_lightning.metrics.functional import accuracy
  from torch import nn
  from torch.nn import functional
  from torch.utils.data import DataLoader, random_split
  from torchvision import transforms
  from torchvision.datasets import MNIST
  from labml import lab, experiment
  from labml.utils.lightning import LabMLLightningLogger
label_desc: Use image node_id 1 to create a new LitMNIST object from inherited base classes: pl with example: obj = LitMNIST(hidden_size, learning_rate)