From b89953bea7dfe6c92164888a66d05bc7d987ef71 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Wed, 11 Jan 2023 22:29:28 +0100 Subject: Fix --- train_ti.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) (limited to 'train_ti.py') diff --git a/train_ti.py b/train_ti.py index e6c437e..951f3dd 100644 --- a/train_ti.py +++ b/train_ti.py @@ -868,11 +868,11 @@ def main(): pass @torch.no_grad() - def on_clip(): + def on_clip(lr): embeddings = text_encoder.text_model.embeddings.temp_token_embedding pre_norm = embeddings.weight.norm(dim=-1, keepdim=True) - lambda_ = min(1.0, 100 * lr_scheduler.get_last_lr()[0]) + lambda_ = min(1.0, 100 * lr) embeddings.weight[:] = F.normalize(embeddings.weight, dim=-1) * (pre_norm + lambda_ * (0.4 - pre_norm)) loop = partial( @@ -991,7 +991,7 @@ def main(): accelerator.backward(loss) if accelerator.sync_gradients: - on_clip() + on_clip(lr_scheduler.get_last_lr()[0]) optimizer.step() if not accelerator.optimizer_step_was_skipped: -- cgit v1.2.3-54-g00ecf