Merge pull request #8 from kohya-ss/dev

Fix text encoder training is not stopped #4. closes #4
kohya-ss · Dec 22, 2022 · d904bb7 · d904bb7
2 parents 689c841 + 0a884da
commit d904bb7
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train_db.py b/train_db.py
@@ -1011,6 +1011,7 @@ def _unscale_grads_replacer(optimizer, inv_scale, found_inf, allow_fp16):
       if stop_text_encoder_training:
         print(f"stop text encoder training at step {global_step}")
         text_encoder.train(False)
+        text_encoder.requires_grad_(False)
 
       with accelerator.accumulate(unet):
         with torch.no_grad():