cuda

app.py

@@ -45,7 +45,7 @@ def process_image_and_text(image, text):
         add_special_tokens=False,
     )
 
-    # inputs = inputs
+    inputs = inputs.to("cuda")
 
     with torch.no_grad():
         generated_ids = model.generate(**inputs, use_cache=True, max_new_tokens=256, do_sample=False)
@@ -76,7 +76,7 @@ if __name__ == "__main__":
     
     # model_path = "/data/shz/project/vlm-r1/VLM-R1/output/Qwen2.5-VL-3B-GRPO-REC/checkpoint-500"
     model_path = "SZhanZ/Qwen2.5VL-VLM-R1-REC-step500"
-    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_path, torch_dtype=torch.float16)
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_path, torch_dtype=torch.float16, device_map="cuda")
     processor = AutoProcessor.from_pretrained(model_path)
     
     def gradio_interface(image, text):