Fix Gradient Checkpointing

2022-12-03 21:24:38 +09:00 · 2022-12-03 21:24:38 +09:00 · 29e7df519b
parent 31dd4f6433
commit 29e7df519b
1 changed files with 2 additions and 3 deletions
--- a/trainer/diffusers_trainer.py
+++ b/trainer/diffusers_trainer.py
@ -741,9 +741,8 @@ def main():

    if args.gradient_checkpointing:
        unet.enable_gradient_checkpointing()
-
-    if args.train_text_encoder:
-        text_encoder.gradient_checkpointing_enable()
+        if args.train_text_encoder:
+            text_encoder.gradient_checkpointing_enable()

    if args.use_xformers:
        unet.set_use_memory_efficient_attention_xformers(True)