code: def batch_augment(x, func, device='/CPU:0'):
    
    with tf.device(device):
      return tf.map_fn(func, x)
docs: 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.