2 files changed, 4 insertions, 4 deletions

diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py
index 0290327..e5e84c8 100644
--- a/training/strategy/dreambooth.py
+++ b/training/strategy/dreambooth.py
@@ -88,8 +88,8 @@ def dreambooth_strategy_callbacks(
 
     def on_prepare():
         unet.requires_grad_(True)
-        text_encoder.requires_grad_(True)
-        text_encoder.text_model.embeddings.requires_grad_(False)
+        text_encoder.text_model.encoder.requires_grad_(True)
+        text_encoder.text_model.final_layer_norm.requires_grad_(True)
 
         if ema_unet is not None:
             ema_unet.to(accelerator.device)
@@ -203,7 +203,7 @@ def dreambooth_prepare(
     lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
     **kwargs
 ):
-    return accelerator.prepare(text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler) + ({})
+    return accelerator.prepare(text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler) + ({},)
 
 
 dreambooth_strategy = TrainingStrategy(
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 732cd74..bd0d178 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -130,7 +130,7 @@ def textual_inversion_strategy_callbacks(
             if lambda_ != 0:
                 w = text_encoder.text_model.embeddings.temp_token_embedding.weight
 
-                mask = torch.zeros(w.size(0), dtype=torch.bool)
+                mask = torch.zeros(w.shape[0], dtype=torch.bool)
                 mask[text_encoder.text_model.embeddings.temp_token_ids] = True
                 mask[zero_ids] = False