benchmarks/model_benchmark/image_classification_benchmark.py/load_model
def load_model():
    model_class = MODEL_MAP[FLAGS.model]
    # Load the EfficientNetV2B0 model and add a classification head.
    model = model_class(include_top=False, weights="imagenet")
    classifier = keras.models.Sequential(
        [
            keras.Input([IMAGE_SIZE[0], IMAGE_SIZE[1], CHANNELS]),
            model,
            keras.layers.GlobalAveragePooling2D(),
            keras.layers.Dense(2),
        ]
    )
    return classifier