Update app.py

app.py

@@ -7,23 +7,22 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 model = torch.hub.load('mair-lab/mapl', 'mapl')
 model.eval()
-model.to(device=device, dtype=torch.bfloat16)
+model.to(device, torch.bfloat16)
 
 
 def predict(image: Image.Image, question: str) -> str:
-    pixel_values = model.image_transform(image).unsqueeze(0).to(device)
+    pixel_values = model.image_transform(image).unsqueeze(0).to(device, torch.bfloat16)
 
     input_ids = None
     if question:
         text = f"Please answer the question. Question: {question} Answer:" if '?' in question else question
         input_ids = model.text_transform(text).input_ids.to(device)
     
-    with torch.autocast(device_type=device, dtype=torch.bfloat16):
-        generated_ids = model.generate(
-            pixel_values=pixel_values,
-            input_ids=input_ids,
-            max_new_tokens=50
-        )
+    generated_ids = model.generate(
+        pixel_values=pixel_values,
+        input_ids=input_ids,
+        max_new_tokens=50
+    )
     
     answer = model.text_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()