Update README.md

README.md

@@ -60,17 +60,19 @@ from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 import requests
 from PIL import Image
 
-model = Pix2StructForConditionalGeneration.from_pretrained('ybelkada/matcha-chartqa').to(0)
+model = Pix2StructForConditionalGeneration.from_pretrained('ybelkada/matcha-chartqa')
 processor = Pix2StructProcessor.from_pretrained('ybelkada/matcha-chartqa')
 url = "https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/val/png/20294671002019.png"
 image = Image.open(requests.get(url, stream=True).raw)
 
-inputs = processor(images=image, text="Is the sum of all 4 places greater than Laos?", return_tensors="pt").to(0)
+inputs = processor(images=image, text="Is the sum of all 4 places greater than Laos?", return_tensors="pt")
 predictions = model.generate(**inputs, max_new_tokens=512)
 print(processor.decode(predictions[0], skip_special_tokens=True))
 >>> No
 ```
 
+To run the predictions on GPU, simply add `.to(0)` when creating the model and when getting the inputs (`inputs = inputs.to(0)`)
+
 # Contribution
 
 This model was originally contributed by Fangyu Liu, Francesco Piccinno et al. and added to the Hugging Face ecosystem by [Younes Belkada](https://huggingface.co/ybelkada).