Merge pull request #411 from ACEsuit/fine_tune_no_swa_bug

Better fix for PR #405, fix patience check when swa is not active
ACEsuit · May 10, 2024 · 61781f4 · 61781f4
2 parents a9681fe + bb10932
commit 61781f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mace/tools/train.py b/mace/tools/train.py
@@ -244,12 +244,12 @@ def train(
 
             if valid_loss >= lowest_loss:
                 patience_counter += 1
-                if patience_counter >= patience and (swa.start is not None and epoch < swa.start):
+                if patience_counter >= patience and (swa is not None and epoch < swa.start):
                     logging.info(
                         f"Stopping optimization after {patience_counter} epochs without improvement and starting swa"
                     )
                     epoch = swa.start
-                elif patience_counter >= patience and (swa.start is None or epoch >= swa.start):
+                elif patience_counter >= patience and (swa is None or epoch >= swa.start):
                     logging.info(
                         f"Stopping optimization after {patience_counter} epochs without improvement"
                     )