torch for gpu usage

app.py

@@ -5,13 +5,14 @@ from huggingface_hub import snapshot_download
 import os
 import cv2
 import tempfile
+import torch
 
 # Function to load YOLO model from a Hugging Face repository
 def load_model(repo_id):
     download_dir = snapshot_download(repo_id)
     print(f"Model downloaded to: {download_dir}")
     model_path = os.path.join(download_dir, "best.pt")  # Ensure best.pt is present in the repository
-    detection_model = YOLO(model_path, device='cuda')
+    detection_model = YOLO(model_path).to('cuda' if torch.cuda.is_available() else 'cpu')  # Move model to appropriate device
     return detection_model
 
 # Load models for traffic cones and license plates

requirements.txt

@@ -2,4 +2,5 @@ ultralytics
 gradio
 huggingface_hub
 Pillow
-opencv-python
+opencv-python
+torch