import tensorflow as tf

  if use_bfloat16:
    preprocessed_resized_image = tf.cast(
      preprocessed_resized_image, tf.bfloat16)
  tensor_dict[fields.InputDataFields.image] = tf.squeeze(
    preprocessed_resized_image, axis=0)
  tensor_dict[fields.InputDataFields.true_image_shape] = tf.squeeze(
