minor fix to (#1170)

mosaicml · May 3, 2024 · bfbb8c5 · bfbb8c5
1 parent c0d591c
commit bfbb8c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llmfoundry/data/utils.py b/llmfoundry/data/utils.py
@@ -136,7 +136,7 @@ def get_num_tokens_in_batch(batch: Batch) -> int:
 def get_text_collator(
     cfg: DictConfig,
     tokenizer: PreTrainedTokenizerBase,
-    dataset_batch_size: int = -1,
+    dataset_batch_size: int,
 ) -> Tuple[Union[transformers.DataCollatorForLanguageModeling,
                  ConcatenatedSequenceCollatorWrapper], int]:
     eos_token_id = cfg.dataset.get('eos_token_id', None)