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()