Merge pull request #405 from ACEsuit/patience_no_swa_specified

Fix check for patience that fails if swa start was not specified
ACEsuit · May 2, 2024 · 6ca53a7 · 6ca53a7
2 parents 1bc9077 + 44e7a2c
commit 6ca53a7
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 epoch < swa.start:
+                if patience_counter >= patience and (swa.start 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 epoch >= swa.start:
+                elif patience_counter >= patience and (swa.start is None or epoch >= swa.start):
                     logging.info(
                         f"Stopping optimization after {patience_counter} epochs without improvement"
                     )