code: def forward_prediction_heads(
      self, output, mask_features, attn_mask_target_size
  ):
      decoder_output = self.decoder_norm(output)
      decoder_output = decoder_output.transpose(0, 1)
  
      if self.task == "caption":
          outputs_caption = (
              decoder_output[:, self.num_queries :] @ self.caping_embed
          )
  
      # recompute class token output.
      norm_decoder_output = decoder_output / (
          decoder_output.norm(dim=-1, keepdim=True) + 1e-7
      )
      obj_token = norm_decoder_output[:, : self.num_queries - 1]
      cls_token = norm_decoder_output[
          :, self.num_queries - 1 : self.num_queries
      ]
  
      sim = (cls_token @ obj_token.transpose(1, 2)).softmax(-1)[
          :, 0, :, None
      ]
      cls_token = (sim * decoder_output[:, : self.num_queries - 1]).sum(
          dim=1, keepdim=True
      )
  
      if self.task == "ref-seg":
          decoder_output = torch.cat(
              (
                  decoder_output[:, : self.num_queries - 1],
                  cls_token,
                  decoder_output[
                      :, self.num_queries : 2 * self.num_queries - 1
                  ],
              ),
              dim=1,
          )
      else:
          decoder_output = torch.cat(
              (decoder_output[:, : self.num_queries - 1], cls_token),
              dim=1,
          )
  
      mask_embed = self.mask_embed(decoder_output)
      outputs_mask = torch.einsum(
          "bqc,bchw->bqhw", mask_embed, mask_features
      )
  
      if is_lower_torch_version():
          attn_mask = F.interpolate(
              outputs_mask,
              size=attn_mask_target_size,
              mode="bicubic",
              align_corners=False,
          )
      else:
          attn_mask = F.interpolate(
              outputs_mask,
              size=attn_mask_target_size,
              mode="bicubic",
              align_corners=False,
              antialias=True,
          )
  
      attn_mask = (
          attn_mask.sigmoid()
          .flatten(2)
          .unsqueeze(1)
          .repeat(1, self.num_heads, 1, 1)
          .flatten(0, 1)
          < 0.5
      ).bool()
      attn_mask = attn_mask.detach()
  
      attn_mask[:, self.num_queries : self.num_queries + 1].fill_(False)
  
      if self.task == "caption":
          results = {
              "attn_mask": attn_mask,
              "outputs_caption": outputs_caption,
          }
          return results
      else:
          class_embed = decoder_output @ self.class_embed
          results = {
              "outputs_mask": outputs_mask,
              "attn_mask": attn_mask,
              "class_embed": class_embed,
          }
          return results
  
file_path: mmdetection/projects/XDecoder/xdecoder/transformer_decoder.py
imports: import torch
  from torch import nn
  from torch.nn import functional
  from mmdet.registry import MODELS
  from .language_model import LanguageEncoder
  from .transformer_blocks import MLP, Conv2d, CrossAttentionLayer, FFNLayer, PositionEmbeddingSine, SelfAttentionLayer
  from .utils import is_lower_torch_version
label_desc: Use image node_id 4 for calling the XDecoderTransformerDecoder obj's underlying member method code with example usage: obj.forward_prediction_heads(output, mask_features, attn_mask_target_size) and returns: results, results