Effectively allow encoder_outputs input to be a tuple in pix2struct (…

…#23932) consistentcy
huggingface · Jun 1, 2023 · 9a35a7b · 9a35a7b
1 parent 9603ef8
commit 9a35a7b
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/transformers/models/pix2struct/modeling_pix2struct.py b/src/transformers/models/pix2struct/modeling_pix2struct.py
@@ -1734,6 +1734,12 @@ def forward(
                 output_hidden_states=output_hidden_states,
                 return_dict=return_dict,
             )
+        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
+            encoder_outputs = BaseModelOutput(
+                last_hidden_state=encoder_outputs[0],
+                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+            )
 
         hidden_states = encoder_outputs[0]