code: def load_data(data_dir):
      """Load training and eval dataset"""
      x, y = np.load(os.path.join(data_dir, "train_data.npy")), np.load(
          os.path.join(data_dir, "train_labels.npy")
      )
  
      vx, vy = np.load(
          os.path.join(data_dir, "eval_data.npy")
      ), np.load(os.path.join(data_dir, "eval_labels.npy"))
  
      print("==== train tensor shape ====")
      print(x.shape, y.shape)
  
      print("==== eval tensor shape ====")
      print(vx.shape, vy.shape)
      # x.shape = (1000, 784), y.shape = (1000, )
  
      x, y = x.astype(np.float32), y.astype(np.int32)
      vx, vy = vx.astype(np.float32), vy.astype(np.int32)
      x /= 255.0
      vx /= 255.0
  
      dtrain = tf.data.Dataset.from_tensor_slices((x, y))
      dtrain = dtrain.map(lambda x, y: (tf.reshape(x, (28, 28, 1)), y))
      dtrain = dtrain.shuffle(10000).batch(512)
  
      deval = tf.data.Dataset.from_tensor_slices((vx, vy))
      deval = deval.map(lambda x, y: (tf.reshape(x, (28, 28, 1)), y))
      deval = deval.batch(10)
      return dtrain, deval
  
file_path: sagemaker-python-sdk/tests/data/tensorflow_mnist/mnist_v2.py
imports: import argparse
  import numpy
  import os
  import json
  import tensorflow
  from tensorflow.keras.layers import Conv2D, BatchNormalization, Dense, Flatten
label_desc: Use image node_id 3 for calling a global function with example usage: load_data(data_dir) and returns: dtrain,  deval