code: def get_dataset():
      # Prepare MNIST dataset by concatenating Train/Test part; we split later.
      train_set = MNIST(
          os.getcwd(),
          download=True,
          transform=transforms.ToTensor(),
          train=True,
      )
      test_set = MNIST(
          os.getcwd(),
          download=True,
          transform=transforms.ToTensor(),
          train=False,
      )
      return train_set, test_set
  
file_path: BentoML/examples/pytorch_mnist/train.py
imports: import argparse
  import os
  import random
  import model
  import numpy
  import torch
  from sklearn.model_selection import KFold
  from torch import nn
  from torchvision import transforms
  from torchvision.datasets import MNIST
  import bentoml
label_desc: Use image node_id 2 for calling a global function with example usage: get_dataset() and returns: train_set,  test_set