1 files changed, 9 insertions, 8 deletions

diff --git a/train_ti.py b/train_ti.py
index 4366c9e..6757bde 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -204,7 +204,8 @@ def parse_args():
         "--vector_shuffle",
         type=str,
         default="auto",
-        help='Vector shuffling algorithm. Choose between ["all", "trailing", "leading", "between", "auto", "off"]',
+        choices=["all", "trailing", "leading", "between", "auto", "off"],
+        help='Vector shuffling algorithm.',
     )
     parser.add_argument(
         "--offset_noise_strength",
@@ -253,10 +254,9 @@ def parse_args():
         "--lr_scheduler",
         type=str,
         default="one_cycle",
-        help=(
-            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
-            ' "constant", "constant_with_warmup", "one_cycle"]'
-        ),
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial",
+                 "constant", "constant_with_warmup", "one_cycle"],
+        help='The scheduler type to use.',
     )
     parser.add_argument(
         "--lr_warmup_epochs",
@@ -280,7 +280,7 @@ def parse_args():
         "--lr_warmup_func",
         type=str,
         default="cos",
-        help='Choose between ["linear", "cos"]'
+        choices=["linear", "cos"],
     )
     parser.add_argument(
         "--lr_warmup_exp",
@@ -292,7 +292,7 @@ def parse_args():
         "--lr_annealing_func",
         type=str,
         default="cos",
-        help='Choose between ["linear", "half_cos", "cos"]'
+        choices=["linear", "half_cos", "cos"],
     )
     parser.add_argument(
         "--lr_annealing_exp",
@@ -330,7 +330,8 @@ def parse_args():
         "--optimizer",
         type=str,
         default="dadan",
-        help='Optimizer to use ["adam", "adam8bit", "lion", "dadam", "dadan", "adafactor"]'
+        choices=["adam", "adam8bit", "lion", "dadam", "dadan", "adafactor"],
+        help='Optimizer to use'
     )
     parser.add_argument(
         "--dadaptation_d0",