fix: Fixed bug when apply yolo model (#22)

rudeuns · web-flow · commit 298e133726f9 · 2024-01-08T22:19:24.000+09:00
#21
diff --git a/yolo/train.py b/yolo/train.py
@@ -2,7 +2,7 @@
 import argparse
 
 def main(opt):
-    model = YOLO()
+    model = YOLO(opt.model)
     model.train(cfg='./custom.yaml', **vars(opt))
 
 def parse_opt():
