StreamingTextDataset takes correct device batch size (#580)

mosaicml · Sep 6, 2023 · 8687a16 · 8687a16
1 parent b8854e1
commit 8687a16
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/llmfoundry/data/text_data.py b/llmfoundry/data/text_data.py
@@ -247,6 +247,7 @@ def build_text_dataloader(
     dataset = StreamingTextDataset(
         tokenizer=tokenizer,
         streams=streams,
+        batch_size=device_batch_size,
         **cfg.dataset,
     )