def preprocess_for_train(image_bytes, use_bfloat16, image_size=IMAGE_SIZE,
                         augment_name=None,
                         randaug_num_layers=None, randaug_magnitude=None):
  """Preprocesses the given image for evaluation.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    use_bfloat16: `bool` for whether to use bfloat16.
    image_size: image size.
    augment_name: `string` that is the name of the augmentation method
      to apply to the image. `autoaugment` if AutoAugment is to be used or
      `randaugment` if RandAugment is to be used. If the value is `None` no
      augmentation method will be applied applied. See autoaugment.py for more
      details.
    randaug_num_layers: 'int', if RandAug is used, what should the number of
      layers be. See autoaugment.py for detailed description.
    randaug_magnitude: 'int', if RandAug is used, what should the magnitude
      be. See autoaugment.py for detailed description.

  Returns:
    A preprocessed image `Tensor`.
  """
  image = _decode_and_random_crop(image_bytes, image_size)
  image = _flip(image)
  image = tf.reshape(image, [image_size, image_size, 3])

  image = tf.image.convert_image_dtype(
      image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)

  if augment_name:
    try:
      import autoaugment  # pylint: disable=g-import-not-at-top
    except ImportError as e:
      logging.exception('Autoaugment is not supported in TF 2.x.')
      raise e

    logging.info('Apply AutoAugment policy %s', augment_name)
    input_image_type = image.dtype
    image = tf.clip_by_value(image, 0.0, 255.0)
    image = tf.cast(image, dtype=tf.uint8)

    if augment_name == 'autoaugment':
      logging.info('Apply AutoAugment policy %s', augment_name)
      image = autoaugment.distort_image_with_autoaugment(image, 'v0')
    elif augment_name == 'randaugment':
      image = autoaugment.distort_image_with_randaugment(
          image, randaug_num_layers, randaug_magnitude)
    else:
      raise ValueError('Invalid value for augment_name: %s' % (augment_name))

    image = tf.cast(image, dtype=input_image_type)
  return image
Preprocesses the given image for evaluation.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    use_bfloat16: `bool` for whether to use bfloat16.
    image_size: image size.
    augment_name: `string` that is the name of the augmentation method
      to apply to the image. `autoaugment` if AutoAugment is to be used or
      `randaugment` if RandAugment is to be used. If the value is `None` no
      augmentation method will be applied applied. See autoaugment.py for more
      details.
    randaug_num_layers: 'int', if RandAug is used, what should the number of
      layers be. See autoaugment.py for detailed description.
    randaug_magnitude: 'int', if RandAug is used, what should the magnitude
      be. See autoaugment.py for detailed description.

  Returns:
    A preprocessed image `Tensor`.
  
tf_to_pytorch/convert_tf_to_pt/original_tf/preprocessing.py