Update c_msel_tr_loss.py

marrlab · Oct 8, 2023 · 23accdb · 23accdb
1 parent 5775c10
commit 23accdb
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/domainlab/algos/msels/c_msel_tr_loss.py b/domainlab/algos/msels/c_msel_tr_loss.py
@@ -17,7 +17,7 @@ def __init__(self, max_es):
         self.max_es = max_es
         super().__init__()
 
-    def update(self):
+    def update(self, clear_counter=False):
         """
         if the best model should be updated
         """
@@ -34,6 +34,9 @@ def update(self):
             logger.info(f"early stop counter: {self.es_c}")
             logger.info(f"loss:{loss}, best loss: {self.best_loss}")
             flag = False  # do not update best model
+            if clear_counter:
+                logger.info("clearing counter")
+                self.es_c = 0
         return flag
 
     def if_stop(self):