def forward(
        self,
        x: torch.Tensor,
        return_attention=False,
        return_qkv=False,
    ):
        """
        Apply the attention mechanism to the input tensor.

        Parameters:
        -----------
        x: torch.Tensor
            Input tensor of shape (Batch, N, C).
        return_attention: bool (default=False)
            If True, return the attention weights.
        return_qkv: bool (default=False)
            If True, return the query, key, and value tensors together with
            the output tensor.
        Returns:
        --------
        x: torch.Tensor
            Output tensor of shape (Batch, N, C).
        qkv: torch.Tensor (optional)
            Query, key, and value tensors of shape
            (Batch, N, 3, num_heads, C // num_heads).
        """
        B, N, _ = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat(
                (
                    self.q_bias,
                    torch.zeros_like(self.v_bias, requires_grad=False),
                    self.v_bias,
                )
            )
        qkv = nn.functional.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = (
            qkv[0],
            qkv[1],
            qkv[2],
        )  # make torchscript happy (cannot use tensor as tuple) (B, H, N, C)
        if self.q_norm is not None:
            q = self.q_norm(q).type_as(v)
        if self.k_norm is not None:
            k = self.k_norm(k).type_as(v)

        q = q * self.scale
        attn = q @ k.transpose(-2, -1)

        if self.relative_position_bias_table is not None:
            relative_position_bias = self.relative_position_bias_table[
                self.relative_position_index.view(-1)
            ].view(
                self.window_size[0] * self.window_size[1] + 1,
                self.window_size[0] * self.window_size[1] + 1,
                -1,
            )  # Wh*Ww,Wh*Ww,nH
            relative_position_bias = relative_position_bias.permute(
                2, 0, 1
            ).contiguous()  # nH, Wh*Ww, Wh*Ww
            attn = attn + relative_position_bias.unsqueeze(0)

        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        if return_attention:
            return attn

        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)

        x = self.proj(x)
        x = self.proj_drop(x)

        if return_qkv:
            return x, qkv

        return x

        Apply the attention mechanism to the input tensor.

        Parameters:
        -----------
        x: torch.Tensor
            Input tensor of shape (Batch, N, C).
        return_attention: bool (default=False)
            If True, return the attention weights.
        return_qkv: bool (default=False)
            If True, return the query, key, and value tensors together with
            the output tensor.
        Returns:
        --------
        x: torch.Tensor
            Output tensor of shape (Batch, N, C).
        qkv: torch.Tensor (optional)
            Query, key, and value tensors of shape
            (Batch, N, 3, num_heads, C // num_heads).
        
braindecode/models/labram.py