code: def __init__(self, args):
      super().__init__()
  
      self.args = args
  
      # mnist images are (1, 28, 28) (channels, height, width)
      self.layer_1 = nn.Linear(28 * 28, 128)
      self.layer_2 = nn.Linear(128, 256)
      self.layer_3 = nn.Linear(256, 10)
  
file_path: guildai/examples/pytorch-lightning/mnist.py
imports: import os
  import argparse
  import torch
  from torch import nn
  from torch.nn import functional
  from torch.optim import Adam
  from torch.utils.data import DataLoader, random_split
  from torchvision import datasets, transforms
  from torchvision.datasets import MNIST
  from pytorch_lightning import Trainer
  from pytorch_lightning.core.lightning import LightningModule
label_desc: Use image node_id 1 to create a new LitMNIST object from inherited base classes: LightningModule with example: obj = LitMNIST(args)