Fix the bug of sd3 controlnet training when using gradient checkpoint…

…ing. (#9498) Fix the bug of sd3 controlnet training when using gradient_checkpointing. Refer to issue #9496
huggingface · Sep 23, 2024 · 00f5b41 · 00f5b41
1 parent 14f6464
commit 00f5b41
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/diffusers/models/controlnet_sd3.py b/src/diffusers/models/controlnet_sd3.py
@@ -336,7 +336,7 @@ def custom_forward(*inputs):
                     return custom_forward
 
                 ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-                hidden_states = torch.utils.checkpoint.checkpoint(
+                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
                     create_custom_forward(block),
                     hidden_states,
                     encoder_hidden_states,