Fix TQDM progress bar cleanup

Lightning-AI · Mar 5, 2024 · 0148c20 · 0148c20
1 parent 1fa6195
commit 0148c20
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/lightning/pytorch/callbacks/progress/tqdm_progress.py b/src/lightning/pytorch/callbacks/progress/tqdm_progress.py
@@ -283,11 +283,12 @@ def on_train_batch_end(
     def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
         if not self.train_progress_bar.disable:
             self.train_progress_bar.set_postfix(self.get_metrics(trainer, pl_module))
+        if self._leave:
+            self.train_progress_bar.close()
 
     @override
     def on_train_end(self, *_: Any) -> None:
-        if not self._leave:
-            self.train_progress_bar.close()
+        self.train_progress_bar.close()
 
     @override
     def on_validation_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: