def _get_attention_mask(
        self, x: Tensor, mask: Optional[Tensor]
    ) -> Optional[Tensor]:
        """Generates an attention mask for causal attention.

        Args:
            x:
                Input tensor of shape (batch_size, sequence_length, channels).
            mask:
                Mask tensor of shape (batch_size, sequence_length) indicating which tokens
                should be masked.

        Returns:
            Attention mask of shape (batch_size, 1, sequence_length, sequence_length).
        """
        B, N = x.shape[:2]

        # Only apply causal attention if mask is not None. This is a bit hacky, but it
        # allows us to use bidirectional instead of causal attention during evaluation
        # and fine-tuning.
        attn_mask = None
        if mask is not None:
            attn_mask = x.new_ones(size=(B, N, N), dtype=torch.bool).tril(diagonal=0)
            # mask has shape (B, N)
            mask = (~mask).unsqueeze(1).expand(B, N, N).bool()
            attn_mask = torch.logical_or(attn_mask, mask)
            attn_mask = attn_mask.unsqueeze(1)  # (B, 1, N, N)
        return attn_mask
Generates an attention mask for causal attention.

        Args:
            x:
                Input tensor of shape (batch_size, sequence_length, channels).
            mask:
                Mask tensor of shape (batch_size, sequence_length) indicating which tokens
                should be masked.

        Returns:
            Attention mask of shape (batch_size, 1, sequence_length, sequence_length).
        
lightly/models/modules/masked_causal_vision_transformer.py