Make sure check for patience does not try to use swa.start if it is None

ACEsuit · May 2, 2024 · 44e7a2c · 44e7a2c
1 parent 1bc9077
commit 44e7a2c
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"
                     )