1 files changed, 4 insertions, 2 deletions

diff --git a/train_ti.py b/train_ti.py
index 507d710..12e3644 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -292,7 +292,7 @@ def parse_args():
     parser.add_argument(
         "--optimizer",
         type=str,
-        default="lion",
+        default="adam",
         help='Optimizer to use ["adam", "adam8bit", "lion"]'
     )
     parser.add_argument(
@@ -586,13 +586,15 @@ def main():
             eps=args.adam_epsilon,
             amsgrad=args.adam_amsgrad,
         )
-    else:
+    elif args.optimizer == 'lion':
         try:
             from lion_pytorch import Lion
         except ImportError:
             raise ImportError("To use Lion, please install the lion_pytorch library: `pip install lion_pytorch`.")
 
         create_optimizer = partial(Lion, use_triton=True)
+    else:
+        raise ValueError(f"Unknown --optimizer \"{args.optimizer}\"")
 
     checkpoint_output_dir = output_dir/"checkpoints"