1 files changed, 5 insertions, 1 deletions

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 9a23a2a..761efbc 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -40,6 +40,8 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         self.position_embedding = embeddings.position_embedding
         self.initializer_factor = config.initializer_factor
 
+        self.decay_target = self.token_embedding.weight[:, :].norm(dim=-1, keepdim=True).median().item()
+
         self.temp_token_embedding = nn.Embedding(
             self.token_embedding.num_embeddings,
             self.token_embedding.embedding_dim,
@@ -99,7 +101,9 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
 
         return embeds
 
-    def normalize(self, target: float = 0.4, lambda_: float = 1.0):
+    def normalize(self, target: Optional[float] = None, lambda_: float = 1.0):
+        if target is None:
+            target = self.decay_target
         w = self.temp_token_embedding.weight
         pre_norm = w[self.temp_token_ids, :].norm(dim=-1, keepdim=True)
         w[self.temp_token_ids] = F.normalize(