From 1e8f37992fa4c0e864d954c1c182ca14880e23e2 Mon Sep 17 00:00:00 2001
From: Francesco Saverio Zuppichini <francesco.zuppichini@gmail.com>
Date: Wed, 9 Mar 2022 15:51:56 +0100
Subject: [PATCH] done (#16012)

---
 .../models/maskformer/modeling_maskformer.py        | 13 ++++++-------
 tests/maskformer/test_modeling_maskformer.py        |  2 +-
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/src/transformers/models/maskformer/modeling_maskformer.py b/src/transformers/models/maskformer/modeling_maskformer.py
index 100d5efe2e..39af8a27eb 100644
--- a/src/transformers/models/maskformer/modeling_maskformer.py
+++ b/src/transformers/models/maskformer/modeling_maskformer.py
@@ -2313,16 +2313,16 @@ class MaskFormerModel(MaskFormerPreTrainedModel):
         )
         queries = transformer_module_output.last_hidden_state
 
+        encoder_hidden_states = None
+        pixel_decoder_hidden_states = None
+        transformer_decoder_hidden_states = None
+        hidden_states = None
+
         if output_hidden_states:
             encoder_hidden_states = pixel_level_module_output.encoder_hidden_states
             pixel_decoder_hidden_states = pixel_level_module_output.decoder_hidden_states
             transformer_decoder_hidden_states = transformer_module_output.hidden_states
             hidden_states = encoder_hidden_states + pixel_decoder_hidden_states + transformer_decoder_hidden_states
-        else:
-            encoder_hidden_states = None
-            pixel_decoder_hidden_states = None
-            transformer_decoder_hidden_states = None
-            hidden_states = None
 
         output = MaskFormerModelOutput(
             encoder_last_hidden_state=image_features,
@@ -2463,7 +2463,6 @@ class MaskFormerForInstanceSegmentation(MaskFormerPreTrainedModel):
         >>> # you can pass them to feature_extractor for postprocessing
         >>> output = feature_extractor.post_process_segmentation(outputs)
         >>> output = feature_extractor.post_process_semantic_segmentation(outputs)
-
         >>> output = feature_extractor.post_process_panoptic_segmentation(outputs)
         ```
         """
@@ -2477,7 +2476,7 @@ class MaskFormerForInstanceSegmentation(MaskFormerPreTrainedModel):
         outputs: MaskFormerModelOutput = self.model(
             pixel_values,
             pixel_mask,
-            output_hidden_states=output_hidden_states,
+            output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss,
             return_dict=True,
             output_attentions=output_attentions,
         )
diff --git a/tests/maskformer/test_modeling_maskformer.py b/tests/maskformer/test_modeling_maskformer.py
index 3f885b3874..2bc3666f1e 100644
--- a/tests/maskformer/test_modeling_maskformer.py
+++ b/tests/maskformer/test_modeling_maskformer.py
@@ -139,7 +139,7 @@ class MaskFormerModelTester:
 
         def comm_check_on_output(result):
             # let's still check that all the required stuff is there
-            self.parent.assertTrue(result.transformer_decoder_hidden_states is not None)
+            self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None)
             self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None)
             self.parent.assertTrue(result.encoder_last_hidden_state is not None)
             # okay, now we need to check the logits shape