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)

    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.
    
examples/mnist/generate_petastorm_mnist.py