From 03137c2032668c216fb53cf8f6acc80e39a505d5 Mon Sep 17 00:00:00 2001
From: vue1999 <vue1999@gmail.com>
Date: Thu, 14 Nov 2024 13:58:01 +0000
Subject: [PATCH] Enable swa if start_swa was given

---
 mace/tools/arg_parser_tools.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mace/tools/arg_parser_tools.py b/mace/tools/arg_parser_tools.py
index da64806a..213ca4e1 100644
--- 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)