diff --git a/README.md b/README.md
index 2bde72c..4378930 100644
--- a/README.md
+++ b/README.md
@@ -41,7 +41,7 @@ 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)
+inputs = feature_extractor(images=image, return_tensors="pt")
 outputs = model(**inputs)
 logits = outputs.logits
 # model predicts one of the 1000 ImageNet classes