def forward(self, x, return_attention=False, return_qkv=False):
        """
        Apply the attention mechanism to the input tensor.
        Parameters
        ----------
        x: torch.Tensor
            Input tensor of shape (Batch, chs, npatchs, patch).
        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
        -------
        torch.Tensor
            Output tensor of shape (Batch, chs, npatchs, patch).
        """

        if return_attention:
            return self.attn(self.norm1(x), return_attention=True)
        if return_qkv:
            y, qkv = self.attn(self.norm1(x), return_qkv=return_qkv)
            x = x + self.drop_path(self.gamma_1 * y)
            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
            return x, qkv

        if self.gamma_1 is None:
            x = x + self.drop_path(self.attn(self.norm1(x)))
            x = x + self.drop_path(self.mlp(self.norm2(x)))
        else:
            x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x)))
            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
        return x

        Apply the attention mechanism to the input tensor.
        Parameters
        ----------
        x: torch.Tensor
            Input tensor of shape (Batch, chs, npatchs, patch).
        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
        -------
        torch.Tensor
            Output tensor of shape (Batch, chs, npatchs, patch).
        
braindecode/models/labram.py