code: def get_dataset():
      def generate_data(_):
          image = tf.ones([500, 500, 3], dtype=tf.float32)
          label = tf.zeros([1], dtype=tf.int32)
          return image, label
  
      def preprocess(image, label):
          label = tf.cast(label, tf.int32)
          label = tf.one_hot(label, NUM_CLASS)
          label = tf.reshape(label, [NUM_CLASS])
          return image, label
  
      dataset = tf.data.Dataset.range(1)
      dataset = dataset.repeat()
      dataset = dataset.map(
          generate_data,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      dataset = dataset.map(
          preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE
      )
      dataset = dataset.repeat()
      dataset = dataset.batch(128, drop_remainder=True)
      return dataset
  
file_path: tensorflow-upstream/tensorflow/python/distribute/integration_test/tpu_memory_test.py
imports: import gc
  import tensorflow
  from tensorflow.python.eager import context
  from tensorflow.python.platform import flags
label_desc: Use image node_id 1 for calling a global function with example usage: get_dataset() and returns: dataset, image,  label, image,  label