code: def setup(self, stage):
      # transform
      transform = transforms.Compose([transforms.ToTensor()])
      mnist_train = MNIST(
          "/tmp/keepsake-test-mnist",
          train=True,
          download=False,
          transform=transform,
      )
      mnist_train = Subset(mnist_train, range(100))
  
      # train/val split
      mnist_train, mnist_val = random_split(mnist_train, [80, 20])
  
      # assign to use in dataloaders
      self.train_dataset = mnist_train
      self.val_dataset = mnist_val
  
imports: import json
  from glob import glob
  import os
  import torch
  from torch.nn import functional
  from torch import nn
  from pytorch_lightning.core.lightning import LightningModule
  from pytorch_lightning import Trainer
  from torch.optim import Adam
  from torch.utils.data import DataLoader, random_split, Subset
  from torch.utils.data.dataset import Dataset
  from torchvision.datasets import MNIST
  from torchvision import transforms
  import pytest
  import urllib
  from keepsake.pl_callback import KeepsakeCallback
label_desc: Use image node_id 4 for calling the ModelNoValidation obj's underlying member method code with example usage: obj.setup(stage) without return types