1 files changed, 4 insertions, 3 deletions

diff --git a/train_lora.py b/train_lora.py
index b8c7396..91bda5c 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -387,7 +387,7 @@ def parse_args():
     parser.add_argument(
         "--optimizer",
         type=str,
-        default="dadan",
+        default="adan",
         choices=["adam", "adam8bit", "adan", "lion", "dadam", "dadan", "adafactor"],
         help='Optimizer to use'
     )
@@ -412,7 +412,7 @@ def parse_args():
     parser.add_argument(
         "--adam_weight_decay",
         type=float,
-        default=1e-2,
+        default=2e-2,
         help="Weight decay to use."
     )
     parser.add_argument(
@@ -780,6 +780,7 @@ def main():
             timm.optim.Adan,
             weight_decay=args.adam_weight_decay,
             eps=args.adam_epsilon,
+            no_prox=True,
         )
     elif args.optimizer == 'lion':
         try:
@@ -961,7 +962,7 @@ def main():
 
         if len(args.placeholder_tokens) != 0:
             params_to_optimize.append({
-                "params": text_encoder.text_model.embeddings.token_override_embedding.parameters(),
+                "params": text_encoder.text_model.embeddings.token_embedding.parameters(),
                 "lr": learning_rate_emb,
                 "weight_decay": 0,
             })