def forward_features(
        self,
        x,
        input_chans=None,
        return_patch_tokens=False,
        return_all_tokens=False,
    ):
        """
        Forward the features of the model.

        Parameters
        ----------
        x : torch.Tensor
            The input data with shape (batch, n_chans, n_patches, patch size),
            if neural decoder or (batch, n_chans, n_times), if neural tokenizer.
        input_chans : int
            The number of input channels.
        return_patch_tokens : bool
            Whether to return the patch tokens.
        return_all_tokens : bool
            Whether to return all the tokens.

        Returns
        -------
        x : torch.Tensor
            The output of the model.
        """

        if self.neural_tokenizer:
            batch_size, nch, n_patch, temporal = self.patch_embed.segment_patch(x).shape
        else:
            batch_size, nch, n_patch = self.patch_embed(x).shape
        x = self.patch_embed(x)
        # add the [CLS] token to the embedded patch tokens
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)

        x = torch.cat((cls_tokens, x), dim=1)

        # Positional Embedding
        if input_chans is not None:
            pos_embed_used = self.position_embedding[:, input_chans]
        else:
            pos_embed_used = self.position_embedding

        if self.position_embedding is not None:
            pos_embed = self._adj_position_embedding(
                pos_embed_used=pos_embed_used, batch_size=batch_size
            )
            x += pos_embed

        # The time embedding is added across the channels after the [CLS] token
        if self.neural_tokenizer:
            num_ch = self.n_chans
        else:
            num_ch = n_patch
        time_embed = self._adj_temporal_embedding(
            num_ch=num_ch, batch_size=batch_size, dim_embed=temporal
        )
        x[:, 1:, :] += time_embed

        x = self.pos_drop(x)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x)
        if self.fc_norm is not None:
            if return_all_tokens:
                return self.fc_norm(x)
            temporal = x[:, 1:, :]
            if return_patch_tokens:
                return self.fc_norm(temporal)
            return self.fc_norm(temporal.mean(1))
        else:
            if return_all_tokens:
                return x
            elif return_patch_tokens:
                return x[:, 1:]
            return x[:, 0]

        Forward the features of the model.

        Parameters
        ----------
        x : torch.Tensor
            The input data with shape (batch, n_chans, n_patches, patch size),
            if neural decoder or (batch, n_chans, n_times), if neural tokenizer.
        input_chans : int
            The number of input channels.
        return_patch_tokens : bool
            Whether to return the patch tokens.
        return_all_tokens : bool
            Whether to return all the tokens.

        Returns
        -------
        x : torch.Tensor
            The output of the model.
        
braindecode/models/labram.py