Fix

1 files changed, 0 insertions, 19 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 49236c6..f0b84b5 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -104,28 +104,10 @@ def textual_inversion_strategy_callbacks(
             yield
 
     @torch.no_grad()
-    def on_before_optimize(epoch: int):
-        if use_emb_decay:
-            params = [
-                p
-                for p in text_encoder.text_model.embeddings.token_embedding.parameters()
-                if p.grad is not None
-            ]
-            return torch.stack(params) if len(params) != 0 else None
-
-    @torch.no_grad()
     def on_after_optimize(w, lrs: dict[str, float]):
         if ema_embeddings is not None:
             ema_embeddings.step(text_encoder.text_model.embeddings.token_embedding.parameters())
 
-        if use_emb_decay and w is not None:
-            lr = lrs["emb"] if "emb" in lrs else lrs["0"]
-            lambda_ = emb_decay * lr
-
-            if lambda_ != 0:
-                norm = w[:, :].norm(dim=-1, keepdim=True)
-                w[:].add_((w[:] / norm.clamp_min(1e-12)) * lambda_ * (emb_decay_target - norm))
-
     def on_log():
         if ema_embeddings is not None:
             return {"ema_decay": ema_embeddings.decay}
@@ -166,7 +148,6 @@ def textual_inversion_strategy_callbacks(
     return TrainingCallbacks(
         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,