Update app.py

app.py

@@ -225,7 +225,6 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
                 padding=True,
                 return_tensors="pt",
             )
-            inputs = inputs.to("cuda")
 
             
         # Otherwise assume it's a LlavaGuard model
@@ -242,7 +241,7 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
             text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
             inputs = processor(text=text_prompt, images=image, return_tensors="pt")
 
-        inputs = {k: v.to('cuda') for k, v in inputs.items()}
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
         with torch.no_grad():
             generated_ids = model.generate(