Use tqdm.auto for notebooks (#298)

Author: hanlint

composer/loggers/tqdm_logger.py

@@ -6,8 +6,8 @@
 from dataclasses import asdict, dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
-import tqdm
 import yaml
+from tqdm import auto
 
 from composer.core.logging import LogLevel, TLogData, TLogDataValue, format_log_data_value
 from composer.core.logging.base_backend import BaseLoggerBackend
@@ -35,7 +35,7 @@ class _TQDMLoggerInstance:
 
     def __init__(self, state: _TQDMLoggerInstanceState) -> None:
         self.state = state
-        self.pbar = tqdm.tqdm(total=state.total,
+        self.pbar = auto.tqdm(total=state.total,
                               desc=state.description,
                               position=state.position,
                               bar_format="{l_bar}{bar:10}{r_bar}{bar:-10b}")
@@ -65,12 +65,12 @@ class TQDMLoggerBackend(BaseLoggerBackend):
 
     Example output::
 
-        Epoch 1: 100%|██████████| 64/64 [00:01<00:00, 53.17it/s, loss/train=2.3023]                                                                                 
-        Epoch 1 (val): 100%|██████████| 20/20 [00:00<00:00, 100.96it/s, accuracy/val=0.0995]  
+        Epoch 1: 100%|██████████| 64/64 [00:01<00:00, 53.17it/s, loss/train=2.3023]
+        Epoch 1 (val): 100%|██████████| 20/20 [00:00<00:00, 100.96it/s, accuracy/val=0.0995]
 
     .. note::
 
-        It is currently not possible to show additional metrics. 
+        It is currently not possible to show additional metrics.
         Custom metrics for the TQDM progress bar will be supported in a future version.
 
     Args:

tests/test_logger.py

@@ -5,8 +5,8 @@
 from unittest.mock import MagicMock
 
 import pytest
-import tqdm
 from _pytest.monkeypatch import MonkeyPatch
+from tqdm import auto
 
 from composer.core.event import Event
 from composer.core.logging import Logger, LogLevel
@@ -82,7 +82,8 @@ def get_mock_tqdm(position: int, *args, **kwargs):
         is_train_to_mock_tqdms[is_train].append(mock_tqdm)
         return mock_tqdm
 
-    monkeypatch.setattr(tqdm, "tqdm", get_mock_tqdm)
+    monkeypatch.setattr(auto, "tqdm", get_mock_tqdm)
+
     max_epochs = 2
     mosaic_trainer_hparams.max_duration = f"{max_epochs}ep"
     mosaic_trainer_hparams.loggers = [TQDMLoggerBackendHparams()]