fix

mosaicml · Jun 8, 2024 · b76b484 · b76b484
1 parent 78998ef
commit b76b484
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llmfoundry/data/dataloader.py b/llmfoundry/data/dataloader.py
@@ -3,7 +3,7 @@
 
 """Dataloader builder utilities."""
 
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 from composer import DataSpec
 from transformers import PreTrainedTokenizerBase
@@ -19,7 +19,7 @@
 def build_dataloader(
     cfg: Dict[str, Any],
     tokenizer: PreTrainedTokenizerBase,
-    device_batch_size: int,
+    device_batch_size: Union[int, float],
 ) -> DataSpec:
     """Builds a dataloader from a config.