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.

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)