1 files changed, 1 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index 0ad7574..a9a2333 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -809,7 +809,7 @@ def main():
             sample_frequency = math.ceil(num_train_epochs * (sample_frequency / args.num_train_steps))
 
         optimizer = create_optimizer(
-            text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
+            text_encoder.text_model.embeddings.token_override_embedding.params.parameters(),
             lr=args.learning_rate,
         )