diff --git a/src/diffusers/models/transformers/transformer_hidream_image.py b/src/diffusers/models/transformers/transformer_hidream_image.py
index 30f6c3a34c9d..f163da3868e3 100644
--- a/src/diffusers/models/transformers/transformer_hidream_image.py
+++ b/src/diffusers/models/transformers/transformer_hidream_image.py
@@ -918,5 +918,5 @@ def forward(
             unscale_lora_layers(self, lora_scale)
 
         if not return_dict:
-            return (output, hidden_states_masks)
-        return Transformer2DModelOutput(sample=output, mask=hidden_states_masks)
+            return (output,)
+        return Transformer2DModelOutput(sample=output)