1 files changed, 2 insertions, 1 deletions

diff --git a/training/ti.py b/training/ti.py
index 1318e22..031fe48 100644
--- a/training/ti.py
+++ b/training/ti.py
@@ -22,7 +22,8 @@ class TrainableEmbeddings(CLIPTextEmbeddings):
         self.train_indices = torch.tensor(new_ids)
 
         self.trainable_embedding = nn.Embedding(self.token_embedding.num_embeddings, self.token_embedding.embedding_dim)
-        self.trainable_embedding.weight.data = self.token_embedding.weight.data.clone()
+        self.trainable_embedding.weight.data.zero_()
+        self.trainable_embedding.weight.data[self.train_indices] = self.token_embedding.weight.data[self.train_indices]
 
     def forward(
         self,