    def forward(self, x_q, x_kv, pad_mask=None, attn_mask=None):
        """
        Forward function.
        Args:
            x_q (`Tensor`):
                Query embeddings.
            x_kv (`Tensor`):
                Key embeddings. Key equals value.
            pad_mask (`int`):
                Padding mask.
            attn_mask (`nn.Module`):
                Attention mask.
        """
        return self.attention(x_q, x_kv, x_kv, key_padding_mask=pad_mask, attn_mask=attn_mask)[0]
Forward function.
Args:
x_q (`Tensor`):
Query embeddings.
x_kv (`Tensor`):
Key embeddings. Key equals value.
pad_mask (`int`):
Padding mask.
attn_mask (`nn.Module`):
Attention mask.
towhee/models/perceiver/multi_head_attention.py