Update train_dreambooth.py

This commit is contained in:
Sourab Mangrulkar
2023-01-15 15:50:25 +01:00
parent 6ddb6c2f8c
commit fd8d3804c7
+2 -2
View File
@@ -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