Fix choice args

1 files changed, 7 insertions, 7 deletions

diff --git a/train_lora.py b/train_lora.py
index a0cd174..538a7f7 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -261,10 +261,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",
@@ -288,7 +287,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",
@@ -300,7 +299,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",
@@ -318,7 +317,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",