1 files changed, 3 insertions, 0 deletions

diff --git a/data/csv.py b/data/csv.py
index 968af8d..dec66d7 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -267,6 +267,9 @@ class VlpnDataModule():
         items = self.prepare_items(template, expansions, items)
         items = self.filter_items(items)
 
+        if (len(items) < self.batch_size):
+            items = (items * self.batch_size)[:self.batch_size]
+
         num_images = len(items)
 
         valid_set_size = min(self.valid_set_size, num_images) if self.valid_set_size is not None else num_images // 10