code: def batch_augment(x, func, device="/CPU:0"):
      """
      Apply dataset augmentation to a batch of exmaples.
      :param x: Tensor representing a batch of examples.
      :param func: Callable implementing dataset augmentation, operating on
        a single image.
      :param device: String specifying which device to use.
      """
      with tf.device(device):
          return tf.map_fn(func, x)
  
label_desc: Use image node_id 2 for calling a global function with example usage: batch_augment(x, func, device) and returns: tf