Enable swa if start_swa was given

vue1999 · Nov 14, 2024 · 03137c2 · 03137c2
1 parent 35c7de6
commit 03137c2
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(
+            (
+                f"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)