test early stop

supervisely-ecosystem · Apr 1, 2024 · 8fd2532 · 8fd2532
1 parent 7bb17d8
commit 8fd2532
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/train/src/main.py b/train/src/main.py
@@ -1302,9 +1302,15 @@ def train_batch_watcher_disable():
 
         threading.Thread(target=train_batch_watcher_func, daemon=True).start()
 
-    def stop_on_batch_end_if_needed(*args, **kwargs):
-        if app.is_stopped():
+    def stop_on_batch_end_if_needed(trainer_validator, *args, **kwargs):
+        try:
+            app_is_stopped = app.is_stopped()
+        except:
+            app_is_stopped = True
+        if app_is_stopped:
+            trainer_validator.stop = True
             raise app.StopException("This error is expected.")
+
 
     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)