Update

1 files changed, 8 insertions, 6 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index ba78b98..e922954 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -117,14 +117,15 @@ def textual_inversion_strategy_callbacks(
         with ema_context():
             yield
 
-    def on_after_optimize(lr: float):
+    @torch.no_grad()
+    def on_before_optimize(lr: float, epoch: int):
         if use_emb_decay:
-            with torch.no_grad():
-                text_encoder.text_model.embeddings.normalize(
-                    emb_decay_target,
-                    min(1.0, emb_decay * lr)
-                )
+            text_encoder.text_model.embeddings.normalize(
+                emb_decay_target,
+                min(1.0, emb_decay * lr)
+            )
 
+    def on_after_optimize(lr: float):
         if ema_embeddings is not None:
             ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters())
 
@@ -154,6 +155,7 @@ def textual_inversion_strategy_callbacks(
         on_model=on_model,
         on_train=on_train,
         on_eval=on_eval,
+        on_before_optimize=on_before_optimize,
         on_after_optimize=on_after_optimize,
         on_log=on_log,
         on_checkpoint=on_checkpoint,