code: def _GetDataSet(batch_size):
      dataset = tfds.load("mnist", split="test")
      dataset = dataset.map(
          map_func=_PreprocessFn, num_parallel_calls=8
      ).batch(batch_size=batch_size)
      dataset = dataset.repeat(count=1)
      return dataset
  
imports: import os.path
  import tempfile
  import tensorflow_datasets
  from tensorflow.compiler.tf2tensorrt._pywrap_py_utils import is_tensorrt_enabled
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.python.client import session
  from tensorflow.python.compiler.tensorrt import trt_convert
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.estimator.estimator import Estimator
  from tensorflow.python.estimator.model_fn import EstimatorSpec
  from tensorflow.python.estimator.model_fn import ModeKeys
  from tensorflow.python.estimator.run_config import RunConfig
  from tensorflow.python.framework import convert_to_constants
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import importer
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.keras.metrics import Accuracy
  from tensorflow.python.layers import layers
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import metrics
  from tensorflow.python.ops import nn
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops.losses import losses
  from tensorflow.python.platform import test
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.saved_model import builder
  from tensorflow.python.saved_model import signature_constants
  from tensorflow.python.saved_model import tag_constants
  from tensorflow.python.saved_model import utils
  from tensorflow.python.saved_model import signature_def_utils
  from tensorflow.python.saved_model.load import load
  from tensorflow.python.summary import summary
  from tensorflow.python.training import saver
  from tensorflow.python.training.adam import AdamOptimizer
  from tensorflow.python.training.checkpoint_management import latest_checkpoint
  from tensorflow.python.training.training_util import get_global_step
label_desc: Use image node_id 2 for calling a global function with example usage: _GetDataSet(batch_size) and returns: dataset