def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        """Forward pass of the attention layer.

        Args:
            x:
                Input tensor of shape (batch_size, sequence_length, channels).
            mask:
                Mask of shape (batch_size, sequence_length) indicating which tokens
                should be masked. Tokens where the mask is True will only be used for
                causal attention, while unmasked tokens are used for bidirectional
                attention. If the mask is None, all tokens are used for bidirectional
                attention.
        """
        B, N, C = x.shape
        attn_mask = self._get_attention_mask(x, mask=mask)
        qkv = (
            self.qkv(x)
            .reshape(B, N, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = qkv.unbind(0)
        q, k = self.q_norm(q), self.k_norm(k)

        if self.fused_attn:
            # Type ignore because only new torch versions support this.
            x = F.scaled_dot_product_attention(  # type: ignore[attr-defined]
                q,
                k,
                v,
                attn_mask=attn_mask,
                dropout_p=self.attn_drop.p if self.training else 0.0,
            )
        else:
            assert False, "Only fused attention is supported for now."
            # TODO: Implement non-fused attention.
            q = q * self.scale  # type: ignore[unreachable]
            attn = q @ k.transpose(-2, -1)
            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
            x = attn @ v

        x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x
Forward pass of the attention layer.

        Args:
            x:
                Input tensor of shape (batch_size, sequence_length, channels).
            mask:
                Mask of shape (batch_size, sequence_length) indicating which tokens
                should be masked. Tokens where the mask is True will only be used for
                causal attention, while unmasked tokens are used for bidirectional
                attention. If the mask is None, all tokens are used for bidirectional
                attention.
        
lightly/models/modules/masked_causal_vision_transformer.py