Fix Issue 782: tqdm_logger did not account for epoch_length argument (#…

…785) * fixes issue #782 * add test case for issue #782 * add test case for issue #782
pytorch · Feb 15, 2020 · 64b8d6f · 64b8d6f
1 parent ae2c49b
commit 64b8d6f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/ignite/contrib/handlers/tqdm_logger.py b/ignite/contrib/handlers/tqdm_logger.py
@@ -218,7 +218,7 @@ def __init__(self, description, metric_names=None, output_transform=None,
     @staticmethod
     def get_max_number_events(event_name, engine):
         if event_name in (Events.ITERATION_STARTED, Events.ITERATION_COMPLETED):
-            return len(engine.state.dataloader)
+            return engine.state.epoch_length
         if event_name in (Events.EPOCH_STARTED, Events.EPOCH_COMPLETED):
             return engine.state.max_epochs
         return 1

diff --git a/tests/ignite/contrib/handlers/test_tqdm_logger.py b/tests/ignite/contrib/handlers/test_tqdm_logger.py
@@ -424,3 +424,19 @@ def test_pbar_on_callable_events(capsys):
     actual = err[-1]
     expected = u'Epoch: [90/100]  90%|█████████  [00:00<00:00]'
     assert actual == expected
+
+
+def test_tqdm_logger_epoch_length(capsys):
+    loader = list(range(100))
+    engine = Engine(update_fn)
+    pbar = ProgressBar(persist=True)
+    pbar.attach(engine)
+    engine.run(loader, epoch_length=50)
+
+    captured = capsys.readouterr()
+    err = captured.err.split('\r')
+    err = list(map(lambda x: x.strip(), err))
+    err = list(filter(None, err))
+    actual = err[-1]
+    expected = u'Epoch: [50/50] 100%|██████████ [00:00<00:00]'
+    assert actual == expected