src/transformers/models/segformer/modeling_tf_segformer.py/TFSegformerDecodeHead/call
class TFSegformerDecodeHead:    def call(self, encoder_hidden_states, training: bool = False):
        batch_size = shape_list(encoder_hidden_states[-1])[0]

        all_hidden_states = ()
        for encoder_hidden_state, mlp in zip(encoder_hidden_states, self.mlps):
            if self.config.reshape_last_stage is False and len(shape_list(encoder_hidden_state)) == 3:
                height = tf.math.sqrt(tf.cast(shape_list(encoder_hidden_state)[1], tf.float32))
                height = width = tf.cast(height, tf.int32)
                encoder_hidden_state = tf.reshape(encoder_hidden_state, (batch_size, height, width, -1))

            # unify channel dimension
            encoder_hidden_state = tf.transpose(encoder_hidden_state, perm=[0, 2, 3, 1])
            height = shape_list(encoder_hidden_state)[1]
            width = shape_list(encoder_hidden_state)[2]
            encoder_hidden_state = mlp(encoder_hidden_state)
            encoder_hidden_state = tf.reshape(encoder_hidden_state, (batch_size, height, width, -1))

            # upsample
            temp_state = tf.transpose(encoder_hidden_states[0], perm=[0, 2, 3, 1])
            upsample_resolution = shape_list(temp_state)[1:-1]
            encoder_hidden_state = tf.image.resize(encoder_hidden_state, size=upsample_resolution, method="bilinear")
            all_hidden_states += (encoder_hidden_state,)

        hidden_states = self.linear_fuse(tf.concat(all_hidden_states[::-1], axis=-1))
        hidden_states = self.batch_norm(hidden_states, training=training)
        hidden_states = self.activation(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)

        # logits of shape (batch_size, height/4, width/4, num_labels)
        logits = self.classifier(hidden_states)

        return logits