From fd8d3804c7441f7c87cdfe96b7d8fc6cdbfd9774 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Sun, 15 Jan 2023 15:50:25 +0100
Subject: [PATCH] Update train_dreambooth.py

---
 examples/lora_dreambooth/train_dreambooth.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/lora_dreambooth/train_dreambooth.py b/examples/lora_dreambooth/train_dreambooth.py
index 50f75ca..2f53123 100644
--- a/examples/lora_dreambooth/train_dreambooth.py
+++ b/examples/lora_dreambooth/train_dreambooth.py
@@ -717,8 +717,8 @@ def main(args):
     if args.gradient_checkpointing:
         unet.enable_gradient_checkpointing()
         # below fails when using lora so commenting it out
-        # if args.train_text_encoder:
-        #     text_encoder.gradient_checkpointing_enable()
+        if args.train_text_encoder and not args.use_lora:
+            text_encoder.gradient_checkpointing_enable()
 
     # Enable TF32 for faster training on Ampere GPUs,
     # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices