code: def mnist_dataset():
      transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,)),
          ]
      )
      data_path = os.path.join(test_data_path(), "mnist")
      os.makedirs(data_path, exist_ok=True)
      with FileLock(os.path.join(data_path, "data.lock")):
          train = datasets.MNIST(
              data_path, train=True, download=True, transform=transform
          )
          test = datasets.MNIST(
              data_path, train=False, download=True, transform=transform
          )
      return train, test
  
file_path: coremltools/coremltools/test/optimize/torch/models/mnist.py
imports: import os
  from collections import OrderedDict
  from coremltools.test.optimize.torch.utils import test_data_path
  import pytest
  import torch.nn
  from filelock import FileLock
  from torchvision import datasets, transforms
label_desc: Use image node_id 4 for calling a global function with example usage: mnist_dataset() and returns: train,  test