LOMO / FIX: Support multiple optimizers (#2745)

huggingface · May 10, 2024 · d8c8634 · d8c8634
1 parent d58a05a
commit d8c8634
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -2035,7 +2035,8 @@ def prepare_optimizer(self, optimizer: torch.optim.Optimizer, device_placement=N
             # transformers & accelerate
             from lomo_optim import AdaLomo, Lomo
 
-            self.has_lomo_optimizer = isinstance(optimizer, (Lomo, AdaLomo))
+            # Support multiple optimizers: https://github.com/huggingface/accelerate/pull/2695#discussion_r1589164607
+            self.has_lomo_optimizer |= isinstance(optimizer, (Lomo, AdaLomo))
 
         # Ensure we can't double wrap an optimizer due to `find_batch_size`
         if getattr(optimizer, "_is_accelerate_prepared", False):