code: def download_mnist_data(download_dir, train=True):
      """
      Downloads the dataset files and returns the torch Dataset object, which
      represents the data as an array of (img, label) pairs.
  
      Each image is a PIL.Image of black-and-white 28x28 pixels.
      Each label is a long integer representing the digit 0..9.
      """
      # This is the only function requiring torch in this module.
  
      # Must import pyarrow before torch. See: https://github.com/uber/petastorm/blob/master/docs/troubleshoot.rst
      import pyarrow  # noqa: F401 pylint: disable=W0611,W0612
      from torchvision import datasets
  
      return datasets.MNIST(
          "{}/{}".format(download_dir, "data"),
          train=train,
          download=True,
      )
  
file_path: petastorm/examples/mnist/generate_petastorm_mnist.py
imports: import argparse
  import numpy
  import os
  import shutil
  import tempfile
  from pyspark.sql import SparkSession
  from examples.mnist import DEFAULT_MNIST_DATA_PATH
  from examples.mnist.schema import MnistSchema
  from petastorm.etl.dataset_metadata import materialize_dataset
  from petastorm.unischema import dict_to_spark_row
label_desc: Use image node_id 2 for calling a global function with example usage: download_mnist_data(download_dir, train) and returns: datasets