1 files changed, 3 insertions, 0 deletions

diff --git a/data/csv.py b/data/csv.py
index c5902ed..913268f 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -2,6 +2,7 @@ import math
 import torch
 import json
 from functools import partial
+import itertools
 from pathlib import Path
 from typing import NamedTuple, Optional, Union, Callable
 
@@ -407,6 +408,7 @@ class VlpnDataset(IterableDataset):
 
             if len(bucket_items) == 0:
                 if len(batch) != 0:
+                    batch = list(itertools.islice(itertools.cycle(batch), batch_size))
                     yield batch
                     batch = []
 
@@ -446,4 +448,5 @@ class VlpnDataset(IterableDataset):
                 batch.append(example)
 
         if len(batch) != 0:
+            batch = list(itertools.islice(itertools.cycle(batch), batch_size))
             yield batch