Update app.py

app.py

@@ -5,15 +5,15 @@ import cv2
 import numpy as np
 from ultralytics import YOLO
 
-# # Clear GPU memory
-# torch.cuda.empty_cache()
+# Clear GPU memory
+torch.cuda.empty_cache()
 
-# # Load YOLO model
-# @st.cache_resource
+# Load YOLO model
+@st.cache_resource
 def load_model():
-    # Dynamically select device (GPU if available, otherwise CPU)
-    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    model = YOLO('./best .pt')  # Replace with your model path
+    Dynamically select device (GPU if available, otherwise CPU)
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model = YOLO('./best .pt').to(device)  # Replace with your model path
     return model
 
 model = load_model()