fix types of num_decays and decay_steps_perc arguments

YerevaNN · Jan 23, 2025 · c3a9417 · c3a9417
1 parent f0608b1
commit c3a9417
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 9 deletions.
diff --git a/submitit_train.py b/submitit_train.py
@@ -51,7 +51,7 @@
                     "--job.config_file",
                     train_config,
                     "--training.steps",
-                    "10000",
+                    "20000",
                 ]
             )
             print(" ".join(function.command))

diff --git a/torchtitan/config_manager.py b/torchtitan/config_manager.py
@@ -230,19 +230,13 @@ def __init__(self):
         )
         self.parser.add_argument(
             "--training.num_decays",
-            type=Optional[float],
+            type=int,
             default=1,
             help="The number of total decays to perform throughout the training, following the WSD-S scheduler",
         )
-        self.parser.add_argument(
-            "--training.decay_steps",
-            type=Optional[int],
-            default=None,
-            help="Steps for lr scheduler decay, default is decay starts immediately after warmup",
-        )
         self.parser.add_argument(
             "--training.decay_steps_perc",
-            type=Optional[float],
+            type=float,
             default=1.0,
             help="The percentage of the steps to use as decay steps",
         )