Merge branch 'develop' into LLPR_universal

SanggyuChong · Sep 17, 2024 · aac749e · aac749e
2 parents ae3b644 + 7a475fb
commit aac749e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 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()

diff --git a/setup.cfg b/setup.cfg
@@ -15,7 +15,7 @@ packages = find:
 python_requires = >=3.7
 install_requires =
     torch>=1.12
-    e3nn==0.4.4
+    e3nn
     numpy<2.0
     opt_einsum
     ase