1 files changed, 2 insertions, 7 deletions

diff --git a/training/util.py b/training/util.py
index 6f42228..2b7f71d 100644
--- a/training/util.py
+++ b/training/util.py
@@ -1,6 +1,7 @@
 from pathlib import Path
 import json
 import copy
+import itertools
 from typing import Iterable, Optional
 from contextlib import contextmanager
 
@@ -71,13 +72,7 @@ class CheckpointerBase:
             file_path = samples_path.joinpath(pool, f"step_{step}.jpg")
             file_path.parent.mkdir(parents=True, exist_ok=True)
 
-            data_enum = enumerate(data)
-
-            batches = [
-                batch
-                for j, batch in data_enum
-                if j * data.batch_size < self.sample_batch_size * self.sample_batches
-            ]
+            batches = list(itertools.islice(itertools.cycle(data), self.sample_batch_size * self.sample_batches))
             prompts = [
                 prompt
                 for batch in batches