Set streaming log level (#1582)

mosaicml · Oct 10, 2024 · 5b0a53e · 5b0a53e
1 parent 813d50e
commit 5b0a53e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/llmfoundry/command_utils/train.py b/llmfoundry/command_utils/train.py
@@ -234,6 +234,9 @@ def train(cfg: DictConfig) -> Trainer:
         logging.getLogger(__name__).setLevel(
             train_cfg.python_log_level.upper(),
         )  # Train script
+        logging.getLogger('streaming').setLevel(
+            train_cfg.python_log_level.upper(),
+        )  # Streaming module
 
     _initialize_dist_with_barrier(dist_timeout=train_cfg.dist_timeout)