diff --git a/hot_fair_utilities/training/yolo_v8_v2/train.py b/hot_fair_utilities/training/yolo_v8_v2/train.py
index d46e2bd..f5ac627 100644
--- a/hot_fair_utilities/training/yolo_v8_v2/train.py
+++ b/hot_fair_utilities/training/yolo_v8_v2/train.py
@@ -102,6 +102,7 @@ def train(
         ),  # Using the environment variable with fallback
         name=name,
         epochs=int(epochs),
+        batch = int(batch_size),
         resume=resume,
         verbose=True,
         deterministic=False,