1 files changed, 8 insertions, 0 deletions

diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index 1f0a117..3f4dbbc 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -138,6 +138,14 @@ def lora_strategy_callbacks(
             state_dict.update(text_encoder_state_dict)
             lora_config["text_encoder_peft_config"] = text_encoder_.get_peft_config_as_dict(inference=True)
 
+            if len(placeholder_tokens) != 0:
+                ti_state_dict = {
+                    f"ti_${token}": text_encoder.text_model.embeddings.get_embed(ids)
+                    for (token, ids)
+                    in zip(placeholder_tokens, placeholder_token_ids)
+                }
+                state_dict.update(ti_state_dict)
+
             save_file(state_dict, checkpoint_output_dir / f"{step}_{postfix}.safetensors")
             with open(checkpoint_output_dir / "lora_config.json", "w") as f:
                 json.dump(lora_config, f)