1 files changed, 1 insertions, 3 deletions

diff --git a/train_ti.py b/train_ti.py
index e933c48..9d06c50 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -405,10 +405,8 @@ class Checkpointer(CheckpointerBase):
         text_encoder = self.accelerator.unwrap_model(self.text_encoder)
 
         for (placeholder_token, placeholder_token_id) in zip(self.placeholder_token, self.placeholder_token_id):
-            training_token_id = self.text_embeddings.id_mapping[placeholder_token_id]
-
             # Save a checkpoint
-            learned_embeds = self.text_embeddings.trainable_embedding.weight[training_token_id]
+            learned_embeds = self.text_embeddings.trainable_embedding.weight[placeholder_token_id]
             learned_embeds_dict = {placeholder_token: learned_embeds.detach().cpu()}
 
             filename = f"%s_%d_%s.bin" % (slugify(placeholder_token), step, postfix)