def gray_augment(is_training=True, **kwargs):
  """Applies random grayscale augmentation."""
  if is_training:
    prob = kwargs['prob'] if 'prob' in kwargs else 0.2
    return [('gray', {'prob': prob})]
  return []