Update main.py

main.py

@@ -24,12 +24,20 @@ from transformers import AutoProcessor, AutoModelForCausalLM
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# main.py (updated YOLO loading)
-from utils import get_yolo_model  # Ensure this import exists
 
-# Load YOLO model correctly
+# main.py (YOLO loading fix)
+from utils import get_yolo_model
+import torch
+
+# Load YOLO model using official method
 yolo_model = get_yolo_model(model_path="weights/icon_detect/best.pt")
-yolo_model.to("cuda" if torch.cuda.is_available() else "cpu")
+
+# Handle device placement
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if str(device) == "cuda":
+    yolo_model = yolo_model.cuda()
+else:
+    yolo_model = yolo_model.cpu()
 
 # Load caption model and processor
 try: