def forward_features(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        """Forward pass of the model without the classification head.

        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.

        Returns:
            Output tensor after applying the transformer blocks.
        """
        x = self.patch_embed(x)
        x = self._pos_embed(x)
        x = self.patch_drop(x)
        x = self.norm_pre(x)
        if self.grad_checkpointing and not jit.is_scripting():
            # TODO: This probably doesn't work correctly as it doesn't consider the
            # mask.
            x = _manipulate.checkpoint_seq(self.blocks, x)
        else:
            for block in self.blocks:
                x = block(x, mask=mask)
        x = self.norm(x)
        return x
Forward pass of the model without the classification head.

        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.

        Returns:
            Output tensor after applying the transformer blocks.
        
lightly/models/modules/masked_causal_vision_transformer.py