1 files changed, 4 insertions, 3 deletions

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 6fda33c..dc4708a 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -37,7 +37,7 @@ def resize_embedding(old_embedding: nn.Embedding, new_num_embeddings: int, initi
 
 
 class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
-    def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings):
+    def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings, dropout_p: float = 0.0):
         super().__init__(config)
 
         self.token_embedding = embeddings.token_embedding
@@ -46,6 +46,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
 
         self.token_override_embedding = PseudoSparseEmbedding(
             self.token_embedding.embedding_dim,
+            dropout_p=dropout_p,
             device=self.token_embedding.weight.device,
             dtype=self.token_embedding.weight.dtype,
         )
@@ -134,7 +135,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         return embeddings
 
 
-def patch_managed_embeddings(text_encoder: CLIPTextModel) -> ManagedCLIPTextEmbeddings:
-    text_embeddings = ManagedCLIPTextEmbeddings(text_encoder.config, text_encoder.text_model.embeddings)
+def patch_managed_embeddings(text_encoder: CLIPTextModel, dropout_p: float = 0.0) -> ManagedCLIPTextEmbeddings:
+    text_embeddings = ManagedCLIPTextEmbeddings(text_encoder.config, text_encoder.text_model.embeddings, dropout_p)
     text_encoder.text_model.embeddings = text_embeddings
     return text_embeddings