1 files changed, 6 insertions, 3 deletions

diff --git a/data/csv.py b/data/csv.py
index 4ebdc1e..480e9f2 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -2,7 +2,6 @@ import math
 import torch
 import json
 from functools import partial
-import itertools
 from pathlib import Path
 from typing import NamedTuple, Optional, Union, Callable
 
@@ -411,7 +410,9 @@ class VlpnDataset(IterableDataset):
             if len(bucket_items) == 0:
                 if len(batch) != 0:
                     if self.fill_batch:
-                        batch = list(itertools.islice(itertools.cycle(batch), batch_size))
+                        fill_items = self.bucket_items[self.bucket_assignments == bucket]
+                        fill_perm = torch.randint(len(fill_items), (batch_size - len(batch),), generator=self.generator)
+                        batch += fill_items[fill_perm]
                     yield batch
                     batch = []
 
@@ -452,5 +453,7 @@ class VlpnDataset(IterableDataset):
 
         if len(batch) != 0:
             if self.fill_batch:
-                batch = list(itertools.islice(itertools.cycle(batch), batch_size))
+                fill_items = self.bucket_items[self.bucket_assignments == bucket]
+                fill_perm = torch.randint(len(fill_items), (batch_size - len(batch),), generator=self.generator)
+                batch += fill_items[fill_perm]
             yield batch