code: def mnist_dataset(data_dir="~/.mnist_qat_data"):
      transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,)),
          ]
      )
      data_path = os.path.expanduser(f"{data_dir}/mnist")
      if not os.path.exists(data_path):
          os.makedirs(data_path)
      train = datasets.MNIST(
          data_path, train=True, download=True, transform=transform
      )
      test = datasets.MNIST(data_path, train=False, transform=transform)
      return train, test
  
file_path: coremltools/docs/examples/optimize/torch/quantization/linear_quantization.py
imports: from collections import OrderedDict
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import os
  from torchvision import datasets, transforms
  from coremltools.optimize.torch.quantization import LinearQuantizer, LinearQuantizerConfig, ModuleLinearQuantizerConfig
  import coremltools
label_desc: Use image node_id 2 for calling a global function with example usage: mnist_dataset(data_dir) and returns: train,  test