Specify the head of the model in the model card

README.md

@@ -17,7 +17,7 @@ widget:
 
 # ONNX convert of ViT (base-sized model)
 
-# Conversion of [ViT-base](https://huggingface.co/google/vit-base-patch16-224)
+Conversion of [ViT-base](https://huggingface.co/google/vit-base-patch16-224), which has a classification head to perform **image classification**.
 
 # Vision Transformer (base-sized model) 
 
@@ -43,25 +43,22 @@ fine-tuned versions on a task that interests you.
 Here is how to use this model to classify an image of the COCO 2017 dataset into one of the 1,000 ImageNet classes:
 
 ```python
-from transformers import ViTFeatureExtractor, ViTForImageClassification
-from PIL import Image
-import requests
-
-url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
-image = Image.open(requests.get(url, stream=True).raw)
-
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
-model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
-
-inputs = feature_extractor(images=image, return_tensors="pt")
-outputs = model(**inputs)
-logits = outputs.logits
-# model predicts one of the 1000 ImageNet classes
-predicted_class_idx = logits.argmax(-1).item()
-print("Predicted class:", model.config.id2label[predicted_class_idx])
-```
+from transformers import AutoFeatureExtractor
+from optimum.onnxruntime import ORTModelForImageClassification
+from optimum.pipelines import pipeline
+
+feature_extractor = AutoFeatureExtractor.from_pretrained("optimum/vit-base-patch16-224")
+# Loading already converted and optimized ORT checkpoint for inference
+model = ORTModelForImageClassification.from_pretrained("optimum/vit-base-patch16-224")
 
-For more code examples, we refer to the [documentation](https://huggingface.co/transformers/model_doc/vit.html#).
+onnx_img_classif = pipeline(
+    "image-classification", model=model, feature_extractor=feature_extractor
+)
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+
+pred = onnx_img_classif(url)
+print("Top-5 predicted classes:", pred)
+```
 
 ## Training data