1 files changed, 2 insertions, 1 deletions

diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py
index bd853e2..3d1abf7 100644
--- a/training/strategy/dreambooth.py
+++ b/training/strategy/dreambooth.py
@@ -98,7 +98,6 @@ def dreambooth_strategy_callbacks(
 
         if cycle < train_text_encoder_cycles:
             text_encoder.train()
-            tokenizer.train()
 
         yield
 
@@ -155,6 +154,8 @@ def dreambooth_strategy_callbacks(
         unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=False)
         text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=False)
 
+        text_encoder_.text_model.embeddings.persist(False)
+
         with ema_context():
             pipeline = VlpnStableDiffusion(
                 text_encoder=text_encoder_,