From 731c1d1cdda9b5dd63430ce9f8e27db6fa22b583 Mon Sep 17 00:00:00 2001 From: TheoLisin Date: Thu, 9 Nov 2023 02:30:51 +0000 Subject: [PATCH] added stop event on validation --- train/src/main.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/train/src/main.py b/train/src/main.py index 9fffacd..f8c4283 100644 --- 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(