def forward(self, image_embedding, image_pe,
                point_embedding) -> Tuple[Tensor, Tensor]:
        """
        Args:
          image_embedding (paddle.Tensor): image to attend to. Should be shape
            B x embedding_dim x h x w for any h and w.
          image_pe (paddle.Tensor): the positional encoding to add to the image. Must
            have the same shape as image_embedding.
          point_embedding (paddle.Tensor): the embedding to add to the query points.
            Must have shape B x N_points x embedding_dim for any N_points.

        Returns:
          paddle.Tensor: the processed point_embedding
          paddle.Tensor: the processed image_embedding
        """
        bs, c, h, w = image_embedding.shape
        image_embedding = image_embedding.flatten(2).transpose([0, 2, 1])
        image_pe = image_pe.flatten(2).transpose([0, 2, 1])
        queries = point_embedding
        keys = image_embedding
        for layer in self.layers:
            queries, keys = layer(
                queries=queries,
                keys=keys,
                query_pe=point_embedding,
                key_pe=image_pe)
        q = queries + point_embedding
        k = keys + image_pe
        attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
        queries = queries + attn_out
        queries = self.norm_final_attn(queries)
        return queries, keys

        Args:
          image_embedding (paddle.Tensor): image to attend to. Should be shape
            B x embedding_dim x h x w for any h and w.
          image_pe (paddle.Tensor): the positional encoding to add to the image. Must
            have the same shape as image_embedding.
          point_embedding (paddle.Tensor): the embedding to add to the query points.
            Must have shape B x N_points x embedding_dim for any N_points.

        Returns:
          paddle.Tensor: the processed point_embedding
          paddle.Tensor: the processed image_embedding
        
contrib/SegmentAnything/segment_anything/modeling/transformer.py