1 files changed, 12 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index 6757bde..fc0d68c 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -330,7 +330,7 @@ def parse_args():
         "--optimizer",
         type=str,
         default="dadan",
-        choices=["adam", "adam8bit", "lion", "dadam", "dadan", "adafactor"],
+        choices=["adam", "adam8bit", "adan", "lion", "dadam", "dadan", "adafactor"],
         help='Optimizer to use'
     )
     parser.add_argument(
@@ -679,6 +679,17 @@ def main():
             eps=args.adam_epsilon,
             amsgrad=args.adam_amsgrad,
         )
+    elif args.optimizer == 'adan':
+        try:
+            import timm.optim
+        except ImportError:
+            raise ImportError("To use Adan, please install the PyTorch Image Models library: `pip install timm`.")
+
+        create_optimizer = partial(
+            timm.optim.Adan,
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+        )
     elif args.optimizer == 'lion':
         try:
             import lion_pytorch