From e6a8c9d269b4952a6944dfe0e78a1f89bd036971 Mon Sep 17 00:00:00 2001
From: Kohya S <ykumeykume@gmail.com>
Date: Thu, 19 Jan 2023 20:39:33 +0900
Subject: [PATCH] Fix some LoRA not trained if gradient checkpointing

---
 train_network.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/train_network.py b/train_network.py
index c0a881ad..03fd01e7 100644
--- a/train_network.py
+++ b/train_network.py
@@ -166,6 +166,9 @@ def train(args):
   if args.gradient_checkpointing:                       # according to TI example in Diffusers, train is required
     unet.train()
     text_encoder.train()
+
+    # set top parameter requires_grad = True for gradient checkpointing works
+    text_encoder.text_model.embeddings.requires_grad_(True)
   else:
     unet.eval()
     text_encoder.eval()