ensure drop_last=True with train data

allenai · May 30, 2023 · ed7c0e8 · ed7c0e8
1 parent 4d986ed
commit ed7c0e8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/olmo/data/__init__.py b/olmo/data/__init__.py
@@ -59,6 +59,7 @@ def build_train_dataloader(train_config: TrainConfig) -> DataLoader:
             max_examples=train_config.global_train_batch_size * train_config.max_duration,
         ),
         batch_size=train_config.device_train_batch_size,
+        drop_last=train_config.data.drop_last,
         collate_fn=collator,
         num_workers=train_config.data.num_workers,
         pin_memory=train_config.data.pin_memory,