1 files changed, 8 insertions, 6 deletions

diff --git a/training/util.py b/training/util.py
index ae6bfc4..60d64f0 100644
--- a/training/util.py
+++ b/training/util.py
@@ -73,20 +73,22 @@ class CheckpointerBase:
             file_path.parent.mkdir(parents=True, exist_ok=True)
 
             batches = list(itertools.islice(itertools.cycle(data), self.sample_batch_size * self.sample_batches))
-            prompts = [
+            prompt_ids = [
                 prompt
                 for batch in batches
-                for prompt in batch["prompts"]
+                for prompt in batch["prompt_ids"]
             ]
-            nprompts = [
+            nprompt_ids = [
                 prompt
                 for batch in batches
-                for prompt in batch["nprompts"]
+                for prompt in batch["nprompt_ids"]
             ]
 
             for i in range(self.sample_batches):
-                prompt = prompts[i * self.sample_batch_size:(i + 1) * self.sample_batch_size]
-                nprompt = nprompts[i * self.sample_batch_size:(i + 1) * self.sample_batch_size]
+                start = i * self.sample_batch_size
+                end = (i + 1) * self.sample_batch_size
+                prompt = prompt_ids[start:end]
+                nprompt = nprompt_ids[start:end]
 
                 samples = pipeline(
                     prompt=prompt,