benchmarks/model_benchmark/image_classification_benchmark.py/load_data
def load_data():
    # Load cats vs dogs dataset, and split into train and validation sets.
    train_dataset, val_dataset = tfds.load(
        "cats_vs_dogs", split=["train[:90%]", "train[90%:]"], as_supervised=True
    )

    resizing = keras.layers.Resizing(
        IMAGE_SIZE[0], IMAGE_SIZE[1], crop_to_aspect_ratio=True
    )

    def preprocess_inputs(image, label):
        image = tf.cast(image, "float32")
        return resizing(image), label

    train_dataset = (
        train_dataset.map(
            preprocess_inputs, num_parallel_calls=tf.data.AUTOTUNE
        )
        .batch(FLAGS.batch_size)
        .prefetch(tf.data.AUTOTUNE)
    )
    val_dataset = (
        val_dataset.map(preprocess_inputs, num_parallel_calls=tf.data.AUTOTUNE)
        .batch(FLAGS.batch_size)
        .cache()
        .prefetch(tf.data.AUTOTUNE)
    )
    return train_dataset, val_dataset