added stop event on validation

supervisely-ecosystem · Nov 9, 2023 · 731c1d1 · 731c1d1
1 parent f2b04f5
commit 731c1d1
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/train/src/main.py b/train/src/main.py
@@ -1147,11 +1147,12 @@ def train_batch_watcher_disable():
 
         threading.Thread(target=train_batch_watcher_func, daemon=True).start()
 
-    def stop_on_train_batch_end_if_needed(*args, **kwargs):
+    def stop_on_batch_end_if_needed(*args, **kwargs):
         if app.app_is_stopped():
             raise app.StopAppError("This error is expected.")
 
-    model.add_callback("on_train_batch_end", stop_on_train_batch_end_if_needed)
+    model.add_callback("on_train_batch_end", stop_on_batch_end_if_needed)
+    model.add_callback("on_val_batch_end", stop_on_batch_end_if_needed)
 
     with app.run_with_stop_app_error_suppression():
         model.train(