diff --git a/src/dataloader_iam.py b/src/dataloader_iam.py
index 5ce2f178226e8e676ea964c4bd1d6760faf7b028..1ef6aa5e3e7c4fdfd4f29fa214e9c438ea2cc0fd 100644
--- a/src/dataloader_iam.py
+++ b/src/dataloader_iam.py
@@ -120,4 +120,4 @@ class DataLoaderIAM:
         gt_texts = [self.samples[i].gt_text for i in batch_range]
 
         self.curr_idx += self.batch_size
-        return Batch(imgs, gt_texts, self.batch_size)
+        return Batch(imgs, gt_texts, len(imgs))