datasets too

mosaicml · Nov 9, 2023 · 242bbe4 · 242bbe4
1 parent eeadca3
commit 242bbe4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/scripts/train/train.py b/scripts/train/train.py
@@ -10,6 +10,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 import torch
+import datasets as hf_datasets
 from composer import Trainer
 from composer.core import Evaluator
 from composer.core.callback import Callback
@@ -39,6 +40,7 @@
 
 transformers_logging.enable_default_handler()
 transformers_logging.disable_progress_bar()
+hf_datasets.logging.disable_progress_bar()
 
 def validate_config(cfg: DictConfig):
     """Validates compatible model and dataloader selection."""