From 35709f91c53e494c29bfe6042dbac35b3f489f10 Mon Sep 17 00:00:00 2001
From: Carolinabanana <140120812+Carolinabanana@users.noreply.github.com>
Date: Thu, 13 Jun 2024 21:14:32 +0100
Subject: [PATCH] Fix gradient checkpointing issue for Stable Diffusion 3

---
 src/diffusers/models/transformers/transformer_sd3.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diffusers/models/transformers/transformer_sd3.py b/src/diffusers/models/transformers/transformer_sd3.py
index 4b159511e25d..740b19bb5377 100644
--- a/src/diffusers/models/transformers/transformer_sd3.py
+++ b/src/diffusers/models/transformers/transformer_sd3.py
@@ -306,7 +306,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,