[update] update the recipe for training

AgriSwarm · Feb 12, 2024 · eec3633 · eec3633
1 parent 43826d1
commit eec3633
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -4,7 +4,7 @@
 
 results = model.train(
     data='custom.yaml',
-    epochs=10, 
+    epochs=30, 
 #    imgsz=640, 
-    device='mps'
+    device='cuda'
 )