code: def load_mnist(img_size):
      compose = transforms.Compose(
          [
              transforms.Resize(img_size),
              transforms.ToTensor(),
              transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
          ]
      )
      output_dir = "./data/mnist"
      return datasets.MNIST(
          root=output_dir, train=True, transform=compose, download=True
      )
  
label_desc: Use image node_id 2 for calling a global function with example usage: load_mnist(img_size) and returns: datasets