Merge pull request #584 from Hongyu-yu/main

Fix bug about undefined swa
ACEsuit · Sep 12, 2024 · d950e18 · d950e18
2 parents 6d768c2 + dcaa0ae
commit d950e18
Showing 1 changed file with 11 additions and 10 deletions.
diff --git a/mace/tools/train.py b/mace/tools/train.py
@@ -279,16 +279,17 @@ def train(
             if rank == 0:
                 if valid_loss >= lowest_loss:
                     patience_counter += 1
-                    if patience_counter >= patience and epoch < swa.start:
-                        logging.info(
-                            f"Stopping optimization after {patience_counter} epochs without improvement and starting Stage Two"
-                        )
-                        epoch = swa.start
-                    elif patience_counter >= patience and epoch >= swa.start:
-                        logging.info(
-                            f"Stopping optimization after {patience_counter} epochs without improvement"
-                        )
-                        break
+                    if patience_counter >= patience:
+                        if swa is not None and epoch < swa.start:
+                            logging.info(
+                                f"Stopping optimization after {patience_counter} epochs without improvement and starting Stage Two"
+                            )
+                            epoch = swa.start
+                        else:
+                            logging.info(
+                                f"Stopping optimization after {patience_counter} epochs without improvement"
+                            )
+                            break
                     if save_all_checkpoints:
                         param_context = (
                             ema.average_parameters()