fix : gradient checkpointing. Added to default as well.

cloneofsimo · cloneofsimo · commit 25aeab489e95 · 2023-01-09T20:14:55.000Z
diff --git a/lora_diffusion/cli_lora_pti.py b/lora_diffusion/cli_lora_pti.py
@@ -598,8 +598,7 @@ def train(
     )
 
     if gradient_checkpointing:
-        text_encoder.gradient_checkpointing_enable()
-        unet.gradient_checkpointing_enable()
+        unet.enable_gradient_checkpointing()
 
     if scale_lr:
         unet_lr = learning_rate_unet * gradient_accumulation_steps * train_batch_size
diff --git a/training_scripts/multivector_example.sh b/training_scripts/multivector_example.sh
@@ -10,6 +10,7 @@ lora_pti \
   --resolution=512 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=4 \
+  --gradient_checkpointing \
   --scale_lr \
   --learning_rate_unet=1e-4 \
   --learning_rate_text=1e-5 \

Original file line number	Diff line number	Diff line change
`@@ -598,8 +598,7 @@ def train(`
`598`	`598`	`)`
`599`	`599`
`600`	`600`	`if gradient_checkpointing:`
`601`		`- text_encoder.gradient_checkpointing_enable()`
`602`		`- unet.gradient_checkpointing_enable()`
	`601`	`+ unet.enable_gradient_checkpointing()`
`603`	`602`
`604`	`603`	`if scale_lr:`
`605`	`604`	`unet_lr = learning_rate_unet * gradient_accumulation_steps * train_batch_size`