diff options
| -rw-r--r-- | training/strategy/lora.py | 5 |
1 files changed, 2 insertions, 3 deletions
diff --git a/training/strategy/lora.py b/training/strategy/lora.py index 89269c0..cfdc504 100644 --- a/training/strategy/lora.py +++ b/training/strategy/lora.py | |||
| @@ -100,9 +100,8 @@ def lora_strategy_callbacks( | |||
| 100 | 100 | ||
| 101 | @torch.no_grad() | 101 | @torch.no_grad() |
| 102 | def on_after_optimize(w, lrs: dict[str, float]): | 102 | def on_after_optimize(w, lrs: dict[str, float]): |
| 103 | lr = lrs["emb"] or lrs["0"] | 103 | if use_emb_decay and w is not None and "emb" in lrs: |
| 104 | 104 | lr = lrs["emb"] | |
| 105 | if use_emb_decay and w is not None: | ||
| 106 | lambda_ = emb_decay * lr | 105 | lambda_ = emb_decay * lr |
| 107 | 106 | ||
| 108 | if lambda_ != 0: | 107 | if lambda_ != 0: |