Merge pull request #697 from vue1999/swa-check

Enable swa if start_swa was given
ACEsuit · Nov 15, 2024 · 3603dcd · 3603dcd
2 parents eca0bb2 + 2efbd5e
commit 3603dcd
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/mace/tools/arg_parser_tools.py b/mace/tools/arg_parser_tools.py
@@ -92,6 +92,15 @@ def check_args(args):
 
     # Loss and optimization
     # Check Stage Two loss start
+    if args.start_swa is not None:
+        args.swa = True
+        log_messages.append(
+            (
+                "Stage Two is activated as start_stage_two was defined",
+                logging.INFO,
+            )
+        )
+
     if args.swa:
         if args.start_swa is None:
             args.start_swa = max(1, args.max_num_epochs // 4 * 3)