1 files changed, 3 insertions, 1 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 33f5fb9..1b5adab 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -64,7 +64,7 @@ def textual_inversion_strategy_callbacks(
     )
 
     def on_accum_model():
-        return text_encoder.text_model.embeddings.overlay
+        return text_encoder.text_model.embeddings
 
     @contextmanager
     def on_train(epoch: int):
@@ -145,6 +145,8 @@ def textual_inversion_prepare(
 
     text_encoder.text_model.encoder.requires_grad_(False)
     text_encoder.text_model.final_layer_norm.requires_grad_(False)
+    text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
+    text_encoder.text_model.embeddings.token_embedding.requires_grad_(False)
     text_encoder.eval()
 
     return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, {}