Fix TI

1 files changed, 2 insertions, 1 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index d735dac..720ebf3 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -114,11 +114,12 @@ def textual_inversion_strategy_callbacks(
             return torch.stack(params) if len(params) != 0 else None
 
     @torch.no_grad()
-    def on_after_optimize(w, lr: float):
+    def on_after_optimize(w, lrs: dict[str, float]):
         if ema_embeddings is not None:
             ema_embeddings.step(text_encoder.text_model.embeddings.token_override_embedding.params.parameters())
 
         if use_emb_decay and w is not None:
+            lr = lrs["emb"] or lrs["0"]
             lambda_ = emb_decay * lr
 
             if lambda_ != 0: