Update app.py

app.py

@@ -183,6 +183,7 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
                 padding=True,
                 return_tensors="pt",
             )
+            inputs.to(model.device)
 
 
         # Otherwise assume it's a LlavaGuard model
@@ -198,11 +199,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")
-
-        processed_inputs = {}
-        for key, value in inputs.items():
-            processed_inputs[key] = value.to(model.device, dtype=torch.bfloat16)
-        inputs = processed_inputs
+            inputs.to(model.device)
 
         with torch.no_grad():
             generated_ids = model.generate(