From ba0b3dbb656011febafe684180b4c8bf53a4bfb1 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 21 Jun 2026 20:01:29 -0600 Subject: [PATCH] Force batch size when bucket is too small by duplicating items in the batch --- toolkit/dataloader_mixins.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index c3da3ff..f878e94 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -199,6 +199,11 @@ class BucketsMixin: for start_idx in range(0, len(bucket.file_list_idx), self.batch_size): end_idx = min(start_idx + self.batch_size, len(bucket.file_list_idx)) batch = bucket.file_list_idx[start_idx:end_idx] + # if the bucket has fewer items left than the requested batch size, + # duplicate items from this batch to pad it up to batch_size + if len(batch) < self.batch_size and len(batch) > 0: + pad = [batch[i % len(batch)] for i in range(self.batch_size - len(batch))] + batch = batch + pad self.batch_indices.append(batch) def shuffle_buckets(self: 'AiToolkitDataset'):