diff --git a/src/eynollah/training/utils.py b/src/eynollah/training/utils.py
index d44f10c..aef67a0 100644
--- a/src/eynollah/training/utils.py
+++ b/src/eynollah/training/utils.py
@@ -1570,7 +1570,7 @@ def data_gen_ocr(
                         if to_yield: yield to_yield
                         
             else:
-                img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y)
+                img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img, batchcount, ret_x, ret_y)
                 if to_yield: yield to_yield