def predict(
        self,
        text_embedding,
        text_embedding_mask,
        encoder_hidden_states,
        encoder_attention_mask,
    ):
        """Generates the bbox coordinates at inference time."""

        merged_encode_hs = torch.cat([encoder_hidden_states, text_embedding],
                                     1)
        merge_att_mask = torch.cat(
            [encoder_attention_mask, text_embedding_mask], 1)

        prefix_ids = torch.repeat_interleave(self.prefix_ids,
                                             merge_att_mask.shape[0],
                                             0).unsqueeze(-1)

        for _ in range(4):
            decoder_output = self.decoder(
                prefix_ids,
                encoder_hidden_states=merged_encode_hs,
                encoder_attention_mask=merge_att_mask,
                labels=None,
                return_dict=True,
            )
            prob_mask = self.bbox_prob_mask.view(1, 1,
                                                 self.bbox_prob_mask.shape[-1])
            prediction_scores = decoder_output.logits + prob_mask

            prefix_ids = torch.cat([
                prefix_ids,
                torch.argmax(prediction_scores[:, -1, :], dim=-1).unsqueeze(1)
            ],
                                   dim=1)

        pred_box = self.process_bbox(prefix_ids[:, 1:])  # xywh 0-1 to xyxy 0-1

        return pred_box
Generates the bbox coordinates at inference time.
mmpretrain/models/heads/grounding_head.py