Update model.py

model.py

@@ -14,16 +14,19 @@ processor = AutoProcessor.from_pretrained(MODEL_ID)
 
 use_cuda = torch.cuda.is_available()
 dtype = torch.bfloat16 if use_cuda else torch.float32
+device = "cuda:0" if use_cuda else "cpu"
 
 model = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID,
-    dtype=dtype,
-    device_map=0 if use_cuda else -1
+    torch_dtype=dtype,
+    device_map=device
 )
 
+model.generation_config.pad_token_id = processor.tokenizer.eos_token_id
 model.eval()
 
 print("MedGemma loaded successfully!")
+print(f"MedGemma loaded on: {device}")
 
 def generate_report(image):
     """