From 7a7725355a71cfeefbd548adb3c90fedae87a58e Mon Sep 17 00:00:00 2001
From: MaxTeselkin <tesmaxxxx@gmail.com>
Date: Tue, 15 Oct 2024 09:45:55 +0000
Subject: [PATCH] fixed progress bar bugs

---
 train/src/main.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/train/src/main.py b/train/src/main.py
index ec21a70..52b3409 100644
--- a/train/src/main.py
+++ b/train/src/main.py
@@ -2676,7 +2676,9 @@ def on_results_file_changed(filepath, pbar):
     watcher = Watcher(
         watch_file,
         on_results_file_changed,
-        progress_bar_epochs(message="Epochs:", total=n_epochs_input.get_value()),
+        progress_bar_epochs(
+            message="Epochs:", total=state.get("n_epochs", n_epochs_input.get_value())
+        ),
     )
     # train model and upload best checkpoints to team files
     device = 0 if torch.cuda.is_available() else "cpu"
@@ -2854,6 +2856,8 @@ def train_model():
     train_thread.start()
     train_thread.join()
     watcher.running = False
+    progress_bar_iters.hide()
+    progress_bar_epochs.hide()
 
     # visualize model predictions
     making_training_vis_f.show()