banao-tech commited on
Commit
3d30079
·
verified ·
1 Parent(s): 85b9952

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +1 -1
main.py CHANGED
@@ -29,7 +29,7 @@ from utils import get_yolo_model
29
  import torch
30
 
31
  # Load YOLO model using official method
32
- yolo_model = get_yolo_model(model_path="weights/icon_detect/best.pt")
33
 
34
  # Handle device placement
35
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
29
  import torch
30
 
31
  # Load YOLO model using official method
32
+ yolo_model = get_yolo_model(model_path="weights/best.pt")
33
 
34
  # Handle device placement
35
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")