refactor: revert some changes

Author: GdoongMathew

src/lightning/pytorch/callbacks/progress/rich_progress.py

@@ -498,6 +498,8 @@ def on_validation_batch_start(
                 visible=False,
             )
 
+        self.refresh()
+
     def _add_task(self, total_batches: Union[int, float], description: str, visible: bool = True) -> "TaskID":
         assert self.progress is not None
         return self.progress.add_task(
@@ -512,22 +514,27 @@ def _initialize_train_progress_bar_id(self) -> None:
         self.train_progress_bar_id = self._add_task(total_batches, train_description)
 
     def _update(
-        self, progress_bar_id: Optional["TaskID"], current: int, visible: bool = True, refresh: bool = True
+        self,
+        progress_bar_id: Optional["TaskID"],
+        current: int,
+        visible: bool = True,
+        hard: bool = False,
     ) -> None:
         if self.progress is not None and self.is_enabled and progress_bar_id is not None:
-            self.progress.update(progress_bar_id, completed=current, visible=visible, refresh=refresh)
+            self.progress.update(progress_bar_id, completed=current, visible=visible)
+            self.refresh(hard=hard)
 
     @override
     def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
         if self.is_enabled and self.val_progress_bar_id is not None and trainer.state.fn == "fit":
             assert self.progress is not None
-            self.progress.update(self.val_progress_bar_id, advance=0, visible=False, refresh=True)
+            self.progress.update(self.val_progress_bar_id, advance=0, visible=False)
+            self.refresh()
 
     @override
     def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
         if trainer.state.fn == "fit":
             self._update_metrics(trainer, pl_module)
-            self.refresh()
         self.reset_dataloader_idx_tracker()
 
     @override
@@ -554,6 +561,7 @@ def on_test_batch_start(
             assert self.progress is not None
             self.progress.update(self.test_progress_bar_id, advance=0, visible=False)
         self.test_progress_bar_id = self._add_task(self.total_test_batches_current_dataloader, self.test_description)
+        self.refresh()
 
     @override
     def on_predict_batch_start(
@@ -573,6 +581,7 @@ def on_predict_batch_start(
         self.predict_progress_bar_id = self._add_task(
             self.total_predict_batches_current_dataloader, self.predict_description
         )
+        self.refresh()
 
     @override
     def on_train_batch_end(
@@ -586,7 +595,7 @@ def on_train_batch_end(
         if not self.is_disabled and self.train_progress_bar_id is None:
             # can happen when resuming from a mid-epoch restart
             self._initialize_train_progress_bar_id()
-        self._update(self.train_progress_bar_id, batch_idx + 1, refresh=False)
+        self._update(self.train_progress_bar_id, batch_idx + 1)
         self._update_metrics(trainer, pl_module)
         self.refresh()
 
@@ -609,12 +618,12 @@ def on_validation_batch_end(
             return
         if trainer.sanity_checking:
             if self.val_sanity_progress_bar_id is not None:
-                self._update(self.val_sanity_progress_bar_id, batch_idx + 1, refresh=True)
+                self._update(self.val_sanity_progress_bar_id, batch_idx + 1)
             return
 
         if self.val_progress_bar_id is None:
             return
-        self._update(self.val_progress_bar_id, batch_idx + 1, refresh=True)
+        self._update(self.val_progress_bar_id, batch_idx + 1)
 
     @override
     def on_test_batch_end(
@@ -628,7 +637,7 @@ def on_test_batch_end(
     ) -> None:
         if self.is_disabled or self.test_progress_bar_id is None:
             return
-        self._update(self.test_progress_bar_id, batch_idx + 1, refresh=True)
+        self._update(self.test_progress_bar_id, batch_idx + 1)
 
     @override
     def on_predict_batch_end(
@@ -642,7 +651,7 @@ def on_predict_batch_end(
     ) -> None:
         if self.is_disabled or self.predict_progress_bar_id is None:
             return
-        self._update(self.predict_progress_bar_id, batch_idx + 1, refresh=True)
+        self._update(self.predict_progress_bar_id, batch_idx + 1)
 
     def _get_train_description(self, current_epoch: int) -> str:
         train_description = f"Epoch {current_epoch}"