import tensorflow as tf

          input_transformation_fn(tensor_dict=tensor_dict))
    self.assertAllClose(transformed_inputs[fields.InputDataFields.image],
                        np_image / 255.)
    self.assertAllClose(transformed_inputs[fields.InputDataFields.
                                           true_image_shape],
                        [4, 4, 3])

  def test_applies_data_augmentation_fn_to_tensor_dict(self):
    np_image = np.random.randint(256, size=(4, 4, 3))
    tensor_dict = {
        fields.InputDataFields.image:
            tf.constant(np_image),
        fields.InputDataFields.groundtruth_classes:
            tf.constant(np.array([3, 1], np.int32))
    }

    def add_one_data_augmentation_fn(tensor_dict):
      return {key: value + 1 for key, value in tensor_dict.items()}

    num_classes = 4
    input_transformation_fn = functools.partial(
        inputs.transform_input_data,
