From ad857ab0df1d9860d095584370b3d1fdb99beae1 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Tue, 4 Apr 2023 09:54:34 +0200
Subject: Fix choice args

---
 train_lora.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

(limited to 'train_lora.py')

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",
-- 
cgit v1.2.3-54-g00ecf