Update README.md

README.md

@@ -25,12 +25,12 @@ model_id = "gokaygokay/paligemma-rich-captions"
 url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
 image = Image.open(requests.get(url, stream=True).raw)
 
-model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).eval()
+model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).to('cuda').eval()
 processor = AutoProcessor.from_pretrained(model_id)
 
 ## prefix
 prompt = "caption en"
-model_inputs = processor(text=prompt, images=image, return_tensors="pt")
+model_inputs = processor(text=prompt, images=image, return_tensors="pt").to('cuda')
 input_len = model_inputs["input_ids"].shape[-1]
 
 with torch.inference_mode():