3 files changed, 2 insertions, 6 deletions

diff --git a/environment.yaml b/environment.yaml
index 42b568f..9c12a0b 100644
--- a/environment.yaml
+++ b/environment.yaml
@@ -11,6 +11,7 @@ dependencies:
     - python=3.10.8
     - pytorch=2.0.0=*cuda11.8*
     - torchvision=0.15.0
+    - xformers=0.0.17.dev481
     - pip:
           - -e .
           - -e git+https://github.com/huggingface/diffusers#egg=diffusers
@@ -24,4 +25,3 @@ dependencies:
           - test-tube>=0.7.5
           - transformers==4.27.1
           - triton==2.0.0
-          - xformers==0.0.17.dev480
diff --git a/train_ti.py b/train_ti.py
index 036c288..7aeff7c 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -620,8 +620,6 @@ def main():
             eps=args.adam_epsilon,
             decouple=True,
         )
-
-        args.learning_rate = 1.0
     elif args.optimizer == 'dadan':
         try:
             import dadaptation
@@ -633,8 +631,6 @@ def main():
             weight_decay=args.adam_weight_decay,
             eps=args.adam_epsilon,
         )
-
-        args.learning_rate = 1.0
     else:
         raise ValueError(f"Unknown --optimizer \"{args.optimizer}\"")
 
diff --git a/training/functional.py b/training/functional.py
index 77f056e..ebb48ab 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -429,7 +429,7 @@ def train_loop(
     try:
         import dadaptation
 
-        isDadaptation = isinstance(optimizer, (dadaptation.DAdaptAdam, dadaptation.DAdaptAdan))
+        isDadaptation = isinstance(optimizer.optimizer, (dadaptation.DAdaptAdam, dadaptation.DAdaptAdan))
     except ImportError:
         pass