Update app.py

app.py

@@ -68,14 +68,14 @@ hf_token = os.getenv("hf_token")
 login(token=hf_token)
 
 processor = AutoProcessor.from_pretrained("google/medgemma-4b-it")
-model = AutoModelForImageTextToText.from_pretrained("google/medgemma-4b-it")
+model = AutoModelForImageTextToText.from_pretrained("google/medgemma-4b-it", device_map = "cpu")
 
 
 model.eval()
 
 # Inference function
 def generate_caption(image, prompt):
-    inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
+    inputs = processor(images=image, text=prompt, return_tensors="pt")
 
     with torch.no_grad():
         outputs = model.generate(