braindecode 1.2.0.dev184328194__py3-none-any.whl → 1.3.0.dev171178473__py3-none-any.whl

braindecode/models/eegnex.py

@@ -16,9 +16,122 @@ from braindecode.modules import Conv2dWithConstraint, LinearWithConstraint
 class EEGNeX(EEGModuleMixin, nn.Module):
     """EEGNeX model from Chen et al. (2024) [eegnex]_.
 
+    :bdg-success:`Convolution`
+
     .. figure:: https://braindecode.org/dev/_static/model/eegnex.jpg
         :align: center
         :alt: EEGNeX Architecture
+        :width: 620px
+
+    .. rubric:: Architectural Overview
+
+    EEGNeX is a **purely convolutional** architecture that refines the EEGNet-style stem
+    and deepens the temporal stack with **dilated temporal convolutions**. The end-to-end
+    flow is:
+
+    - (i) **Block-1/2**: two temporal convolutions ``(1 x L)`` with BN refine a
+      learned FIR-like *temporal filter bank* (no pooling yet);
+    - (ii) **Block-3**: depthwise **spatial** convolution across electrodes
+      ``(n_chans x 1)`` with max-norm constraint, followed by ELU → AvgPool (time) → Dropout;
+    - (iii) **Block-4/5**: two additional **temporal** convolutions with increasing **dilation**
+      to expand the receptive field; the last block applies ELU → AvgPool → Dropout → Flatten;
+    - (iv) **Classifier**: a max-norm–constrained linear layer.
+
+    The published work positions EEGNeX as a compact, conv-only alternative that consistently
+    outperforms prior baselines across MOABB-style benchmarks, with the popular
+    “EEGNeX-8,32” shorthand denoting *8 temporal filters* and *kernel length 32*.
+
+
+    .. rubric:: Macro Components
+
+    - **Block-1 / Block-2 — Temporal filter (learned).**
+
+       - *Operations.*
+       - :class:`torch.nn.Conv2d` with kernels ``(1, L)``
+       - :class:`torch.nn.BatchNorm2d` (no nonlinearity until Block-3, mirroring a linear FIR analysis stage).
+         These layers set up frequency-selective detectors before spatial mixing.
+
+       - *Interpretability.* Kernels can be inspected as FIR filters; two stacked temporal
+         convs allow longer effective kernels without parameter blow-up.
+
+    - **Block-3 — Spatial projection + condensation.**
+
+        - *Operations.*
+        - :class:`braindecode.modules.Conv2dWithConstraint` with kernel``(n_chans, 1)``
+          and ``groups = filter_2`` (depthwise across filters)
+        - :class:`torch.nn.BatchNorm2d`
+        - :class:`torch.nn.ELU`
+        - :class:`torch.nn.AvgPool2d` (time)
+        - :class:`torch.nn.Dropout`.
+
+    **Role**: Learns per-filter spatial patterns over the **full montage** while temporal
+      pooling stabilizes and compresses features; max-norm encourages well-behaved spatial
+      weights similar to EEGNet practice.
+
+    - **Block-4 / Block-5 — Dilated temporal integration.**
+
+        - *Operations.*
+        - :class:`torch.nn.Conv2d` with kernels ``(1, k)`` and **dilations**
+          (e.g., 2 then 4);
+        - :class:`torch.nn.BatchNorm2d`
+        - :class:`torch.nn.ELU`
+        - :class:`torch.nn.AvgPool2d` (time)
+        - :class:`torch.nn.Dropout`
+        - :class:`torch.nn.Flatten`.
+
+    **Role**: Expands the temporal receptive field efficiently to capture rhythms and
+    long-range context after condensation.
+
+    - **Final Classifier — Max-norm linear.**
+
+        - *Operations.*
+        - :class:`braindecode.modules.LinearWithConstraint` maps the flattened
+          vector to the target classes; the max-norm constraint regularizes the readout.
+
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal (where time-domain patterns are learned).**
+      Blocks 1-2 learn the primary filter bank (oscillations/transients), while Blocks 4-5
+      use **dilation** to integrate over longer horizons without extra pooling. The final
+      AvgPool in Block-5 sets the output token rate and helps noise suppression.
+
+    - **Spatial (how electrodes are processed).**
+      A *single* depthwise spatial conv (Block-3) spans the entire electrode set
+      (kernel ``(n_chans, 1)``), producing per-temporal-filter topographies; no cross-filter
+      mixing occurs at this stage, aiding interpretability.
+
+    - **Spectral (how frequency content is captured).**
+      Frequency selectivity emerges from the learned temporal kernels; dilation broadens effective
+      bandwidth coverage by composing multiple scales.
+
+    .. rubric:: Additional Mechanisms
+
+    - **EEGNeX-8,32 naming.** “8,32” indicates *8 temporal filters* and *kernel length 32*,
+      reflecting the paper's ablation path from EEGNet-8,2 toward thicker temporal kernels
+      and a deeper conv stack.
+    - **Max-norm constraints.** Spatial (Block-3) and final linear layers use max-norm
+      regularization—standard in EEG CNNs—to reduce overfitting and encourage stable spatial
+      patterns.
+
+    .. rubric:: Usage and Configuration
+
+    - **Kernel schedule.** Start with the canonical **EEGNeX-8,32** (``filter_1=8``,
+      ``kernel_block_1_2=32``) and keep **Block-3** depth multiplier modest (e.g., 2) to match
+      the paper's “pure conv” profile.
+    - **Pooling vs. dilation.** Use pooling in Blocks 3 and 5 to control compute and variance;
+      increase dilations (Blocks 4-5) to widen temporal context when windows are short.
+    - **Regularization.** Combine dropout (Blocks 3 & 5) with max-norm on spatial and
+      classifier layers; prefer ELU activations for stable training on small EEG datasets.
+
+
+    - The braindecode implementation follows the paper's conv-only design with five blocks
+      and reproduces the depthwise spatial step and dilated temporal stack. See the class
+      reference for exact kernel sizes, dilations, and pooling defaults. You can check the
+      original implementation at [EEGNexCode]_.
+
+    .. versionadded:: 1.1
+
 
     Parameters
     ----------
@@ -45,12 +158,6 @@ class EEGNeX(EEGModuleMixin, nn.Module):
     avg_pool_block5 : tuple[int, int], optional
         Pooling size for block 5. Default is (1, 8).
 
-    Notes
-    -----
-    This implementation is not guaranteed to be correct, has not been checked
-    by original authors, only reimplemented from the paper description and
-    source code in tensorflow [EEGNexCode]_.
-
     References
     ----------
     .. [eegnex] Chen, X., Teng, X., Chen, H., Pan, Y., & Geyer, P. (2024).

braindecode/models/eegsimpleconv.py

@@ -21,6 +21,8 @@ from braindecode.models.base import EEGModuleMixin
 class EEGSimpleConv(EEGModuleMixin, torch.nn.Module):
     """EEGSimpleConv from Ouahidi, YE et al. (2023) [Yassine2023]_.
 
+    :bdg-success:`Convolution`
+
     .. figure:: https://raw.githubusercontent.com/elouayas/EEGSimpleConv/refs/heads/main/architecture.png
         :align: center
         :alt: EEGSimpleConv Architecture

braindecode/models/eegtcnet.py

@@ -157,7 +157,7 @@ class EEGTCNet(EEGModuleMixin, nn.Module):
 class _EEGNetTC(nn.Module):
     """EEGNet Temporal Convolutional Network (TCN) block.
 
-    The main difference from our EEGNetV4 (braindecode) implementation is the
+    The main difference from our :class:`EEGNet` (braindecode) implementation is the
     kernel and dimensional order. Because of this, we decided to keep this
     implementation in a future issue; we will re-evaluate if it is necessary
     to maintain this separate implementation.

braindecode/models/labram.py

@@ -22,6 +22,8 @@ from braindecode.modules import MLP, DropPath
 class Labram(EEGModuleMixin, nn.Module):
     """Labram from Jiang, W B et al (2024) [Jiang2024]_.
 
+    :bdg-danger:`Large Brain Model`
+
     .. figure:: https://arxiv.org/html/2405.18765v1/x1.png
         :align: center
         :alt: Labram Architecture.
@@ -43,31 +45,45 @@ class Labram(EEGModuleMixin, nn.Module):
     equals True. The original implementation uses (batch, n_chans, n_patches,
     patch_size) as input with static segmentation of the input data.
 
-    The models have the following sequence of steps:
-    if neural tokenizer:
-        - SegmentPatch: Segment the input data in patches;
-        - TemporalConv: Apply a temporal convolution to the segmented data;
-        - Residual adding cls, temporal and position embeddings (optional);
-        - WindowsAttentionBlock: Apply a windows attention block to the data;
-        - LayerNorm: Apply layer normalization to the data;
-        - Linear: An head linear layer to transformer the data into classes.
-
-    else:
-        - PatchEmbed: Apply a patch embedding to the input data;
-        - Residual adding cls, temporal and position embeddings (optional);
-        - WindowsAttentionBlock: Apply a windows attention block to the data;
-        - LayerNorm: Apply layer normalization to the data;
-        - Linear: An head linear layer to transformer the data into classes.
+    The models have the following sequence of steps::
+
+        if neural tokenizer:
+            - SegmentPatch: Segment the input data in patches;
+            - TemporalConv: Apply a temporal convolution to the segmented data;
+            - Residual adding cls, temporal and position embeddings (optional);
+            - WindowsAttentionBlock: Apply a windows attention block to the data;
+            - LayerNorm: Apply layer normalization to the data;
+            - Linear: An head linear layer to transformer the data into classes.
+
+        else:
+            - PatchEmbed: Apply a patch embedding to the input data;
+            - Residual adding cls, temporal and position embeddings (optional);
+            - WindowsAttentionBlock: Apply a windows attention block to the data;
+            - LayerNorm: Apply layer normalization to the data;
+            - Linear: An head linear layer to transformer the data into classes.
 
     .. versionadded:: 0.9
 
+
+    Examples
+    --------
+    Load pre-trained weights::
+
+        >>> import torch
+        >>> from braindecode.models import Labram
+        >>> model = Labram(n_times=1600, n_chans=64, n_outputs=4)
+        >>> url = "https://huggingface.co/braindecode/Labram-Braindecode/blob/main/braindecode_labram_base.pt"
+        >>> state = torch.hub.load_state_dict_from_url(url, progress=True)
+        >>> model.load_state_dict(state)
+
+
     Parameters
     ----------
     patch_size : int
         The size of the patch to be used in the patch embedding.
     emb_size : int
         The dimension of the embedding.
-    in_channels : int
+    in_conv_channels : int
         The number of convolutional input channels.
     out_channels : int
         The number of convolutional output channels.
@@ -79,8 +95,10 @@ class Labram(EEGModuleMixin, nn.Module):
         The expansion ratio of the mlp layer
     qkv_bias :  bool (default=False)
         If True, add a learnable bias to the query, key, and value tensors.
-    qk_norm : Pytorch Normalize layer (default=None)
-        If not None, apply LayerNorm to the query and key tensors
+    qk_norm : Pytorch Normalize layer (default=nn.LayerNorm)
+        If not None, apply LayerNorm to the query and key tensors.
+        Default is nn.LayerNorm for better weight transfer from original LaBraM.
+        Set to None to disable Q,K normalization.
     qk_scale : float (default=None)
         If not None, use this value as the scale factor. If None,
         use head_dim**-0.5, where head_dim = dim // num_heads.
@@ -92,9 +110,10 @@ class Labram(EEGModuleMixin, nn.Module):
         Dropout rate for the attention weights used on DropPath.
     norm_layer : Pytorch Normalize layer (default=nn.LayerNorm)
         The normalization layer to be used.
-    init_values : float (default=None)
+    init_values : float (default=0.1)
         If not None, use this value to initialize the gamma_1 and gamma_2
-        parameters.
+        parameters for residual scaling. Default is 0.1 for better weight
+        transfer from original LaBraM. Set to None to disable.
     use_abs_pos_emb : bool (default=True)
         If True, use absolute position embedding.
     use_mean_pooling : bool (default=True)
@@ -135,19 +154,19 @@ class Labram(EEGModuleMixin, nn.Module):
         input_window_seconds=None,
         patch_size=200,
         emb_size=200,
-        in_channels=1,
+        in_conv_channels=1,
         out_channels=8,
         n_layers=12,
         att_num_heads=10,
         mlp_ratio=4.0,
         qkv_bias=False,
-        qk_norm=None,
+        qk_norm=nn.LayerNorm,
         qk_scale=None,
         drop_prob=0.0,
         attn_drop_prob=0.0,
         drop_path_prob=0.0,
         norm_layer=nn.LayerNorm,
-        init_values=None,
+        init_values=0.1,
         use_abs_pos_emb=True,
         use_mean_pooling=True,
         init_scale=0.001,
@@ -183,15 +202,15 @@ class Labram(EEGModuleMixin, nn.Module):
             self.patch_size = patch_size
         self.n_path = self.n_times // self.patch_size
 
-        if neural_tokenizer and in_channels != 1:
+        if neural_tokenizer and in_conv_channels != 1:
             warn(
                 "The model is in Neural Tokenizer mode, but the variable "
-                + "`in_channels` is different from the default values."
-                + "`in_channels` is only needed for the Neural Decoder mode."
-                + "in_channels is not used in the Neural Tokenizer mode.",
+                + "`in_conv_channels` is different from the default values."
+                + "`in_conv_channels` is only needed for the Neural Decoder mode."
+                + "in_conv_channels is not used in the Neural Tokenizer mode.",
                 UserWarning,
             )
-            in_channels = 1
+            in_conv_channels = 1
             # If you can use the model in Neural Tokenizer mode,
         # temporal conv layer will be use over the patched dataset
         if neural_tokenizer:
@@ -228,7 +247,7 @@ class Labram(EEGModuleMixin, nn.Module):
                 _PatchEmbed(
                     n_times=self.n_times,
                     patch_size=patch_size,
-                    in_channels=in_channels,
+                    in_channels=in_conv_channels,
                     emb_dim=self.emb_size,
                 ),
             )
@@ -373,8 +392,7 @@ class Labram(EEGModuleMixin, nn.Module):
         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.
+            The input data with shape (batch, n_chans, n_times).
         input_chans : int
             The number of input channels.
         return_patch_tokens : bool
@@ -387,37 +405,72 @@ class Labram(EEGModuleMixin, nn.Module):
         x : torch.Tensor
             The output of the model.
         """
+        batch_size = x.shape[0]
+
         if self.neural_tokenizer:
-            batch_size, nch, n_patch, temporal = self.patch_embed.segment_patch(x).shape
+            # For neural tokenizer: input is (batch, n_chans, n_times)
+            # patch_embed returns (batch, n_chans, emb_dim)
+            x = self.patch_embed(x)
+            # x shape: (batch, n_chans, emb_dim)
+            n_patch = self.n_chans
+            temporal = self.emb_size
         else:
-            batch_size, nch, n_patch = self.patch_embed(x).shape
-        x = self.patch_embed(x)
+            # For neural decoder: input is (batch, n_chans, n_times)
+            # patch_embed returns (batch, n_patchs, emb_dim)
+            x = self.patch_embed(x)
+            # x shape: (batch, n_patchs, emb_dim)
+            batch_size, n_patch, temporal = x.shape
+
         # add the [CLS] token to the embedded patch tokens
         cls_tokens = self.cls_token.expand(batch_size, -1, -1)
 
+        # Concatenate cls token with patch/channel embeddings
         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
-            )
+            if self.neural_tokenizer:
+                # In tokenizer mode, use channel-based position embedding
+                if input_chans is not None:
+                    pos_embed_used = self.position_embedding[:, input_chans]
+                else:
+                    pos_embed_used = self.position_embedding
+
+                pos_embed = self._adj_position_embedding(
+                    pos_embed_used=pos_embed_used, batch_size=batch_size
+                )
+            else:
+                # In decoder mode, we have different number of patches
+                # Adapt position embedding for n_patch patches
+                # Use the first n_patch+1 positions from position_embedding
+                n_pos = min(self.position_embedding.shape[1], n_patch + 1)
+                pos_embed_used = self.position_embedding[:, :n_pos, :]
+                pos_embed = pos_embed_used.expand(batch_size, -1, -1)
+
             x += pos_embed
 
         # The time embedding is added across the channels after the [CLS] token
         if self.neural_tokenizer:
             num_ch = self.n_chans
+            time_embed = self._adj_temporal_embedding(
+                num_ch=num_ch, batch_size=batch_size, dim_embed=temporal
+            )
+            x[:, 1:, :] += time_embed
         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
+            # In decoder mode, we have n_patch patches and don't need to expand
+            # Just broadcast the temporal embedding
+            if temporal is None:
+                temporal = self.emb_size
+
+            # Get temporal embeddings for n_patch patches
+            n_time_tokens = min(n_patch, self.temporal_embedding.shape[1] - 1)
+            time_embed = self.temporal_embedding[
+                :, 1 : n_time_tokens + 1, :
+            ]  # (1, n_patch, emb_dim)
+            time_embed = time_embed.expand(
+                batch_size, -1, -1
+            )  # (batch, n_patch, emb_dim)
+            x[:, 1:, :] += time_embed
 
         x = self.pos_drop(x)
 
@@ -428,10 +481,10 @@ class Labram(EEGModuleMixin, nn.Module):
         if self.fc_norm is not None:
             if return_all_tokens:
                 return self.fc_norm(x)
-            temporal = x[:, 1:, :]
+            tokens = x[:, 1:, :]
             if return_patch_tokens:
-                return self.fc_norm(temporal)
-            return self.fc_norm(temporal.mean(1))
+                return self.fc_norm(tokens)
+            return self.fc_norm(tokens.mean(1))
         else:
             if return_all_tokens:
                 return x
@@ -505,14 +558,16 @@ class Labram(EEGModuleMixin, nn.Module):
     def _adj_temporal_embedding(self, num_ch, batch_size, dim_embed=None):
         """
         Adjust the dimensions of the time embedding to match the
-        number of channels.
+        number of channels or patches.
 
         Parameters
         ----------
         num_ch : int
-            The number of channels or number of code books vectors.
+            The number of channels or number of patches.
         batch_size : int
             Batch size of the input data.
+        dim_embed : int
+            The embedding dimension (temporal feature dimension).
 
         Returns
         -------
@@ -523,17 +578,24 @@ class Labram(EEGModuleMixin, nn.Module):
         if dim_embed is None:
             cut_dimension = self.patch_size
         else:
-            cut_dimension = dim_embed
-        # first step will be match the time_embed to the number of channels
-        temporal_embedding = self.temporal_embedding[:, 1:cut_dimension, :]
+            cut_dimension = min(dim_embed, self.temporal_embedding.shape[1] - 1)
+
+        # Get the temporal embedding: (1, temporal_embedding_dim, emb_size)
+        # Slice to cut_dimension: (1, cut_dimension, emb_size)
+        temporal_embedding = self.temporal_embedding[:, 1 : cut_dimension + 1, :]
+
         # Add a new dimension to the time embedding
-        # e.g. (batch, 62, 200) -> (batch, 1, 62, 200)
+        # e.g. (1, 5, 200) -> (1, 1, 5, 200)
         temporal_embedding = temporal_embedding.unsqueeze(1)
-        # Expand the time embedding to match the number of channels
-        # or number of patches from
+
+        # Expand the time embedding to match the number of channels or patches
+        # (1, 1, cut_dimension, 200) -> (batch_size, num_ch, cut_dimension, 200)
         temporal_embedding = temporal_embedding.expand(batch_size, num_ch, -1, -1)
+
         # Flatten the intermediate dimensions
+        # (batch_size, num_ch, cut_dimension, 200) -> (batch_size, num_ch * cut_dimension, 200)
         temporal_embedding = temporal_embedding.flatten(1, 2)
+
         return temporal_embedding
 
     def _adj_position_embedding(self, pos_embed_used, batch_size):
@@ -679,25 +741,27 @@ class _SegmentPatch(nn.Module):
 
 
 class _PatchEmbed(nn.Module):
-    """EEG to Patch Embedding.
+    """EEG to Patch Embedding for Neural Decoder mode.
 
     This code is used when we want to apply the patch embedding
-    after the codebook layer.
+    after the codebook layer (Neural Decoder mode).
+
+    The input is expected to be in the format (Batch, n_channels, n_times),
+    but the original LaBraM expects pre-patched data (Batch, n_channels, n_patches, patch_size).
+    This class reshapes the input to the pre-patched format, then applies a 2D
+    convolution to project this pre-patched data to the embedding dimension,
+    and finally flattens across channels to produce a unified embedding.
 
     Parameters:
     -----------
     n_times: int (default=2000)
-        Number of temporal components of the input tensor.
+        Number of temporal components of the input tensor (used for dimension calculation).
     patch_size: int (default=200)
         Size of the patch, default is 1-seconds with 200Hz.
     in_channels: int (default=1)
-        Number of input channels for to be used in the convolution.
+        Number of input channels (from VQVAE codebook).
     emb_dim: int (default=200)
-        Number of out_channes to be used in the convolution, here,
-        we used the same as patch_size.
-    n_codebooks: int (default=62)
-        Number of patches to be used in the convolution, here,
-        we used the same as n_times // patch_size.
+        Number of output embedding dimension.
     """
 
     def __init__(
@@ -707,10 +771,13 @@ class _PatchEmbed(nn.Module):
         self.n_times = n_times
         self.patch_size = patch_size
         self.patch_shape = (1, self.n_times // self.patch_size)
-        n_patchs = n_codebooks * (self.n_times // self.patch_size)
-
-        self.n_patchs = n_patchs
+        self.n_patchs = self.n_times // self.patch_size
+        self.emb_dim = emb_dim
+        self.in_channels = in_channels
 
+        # 2D Conv to project the pre-patched data
+        # Input: (Batch, in_channels, n_patches, patch_size)
+        # After proj: (Batch, emb_dim, n_patches, 1)
         self.proj = nn.Conv2d(
             in_channels=in_channels,
             out_channels=emb_dim,
@@ -718,27 +785,64 @@ class _PatchEmbed(nn.Module):
             stride=(1, self.patch_size),
         )
 
-        self.merge_transpose = Rearrange(
-            "Batch ch patch spatch -> Batch patch spatch ch",
-        )
-
     def forward(self, x):
         """
-        Apply the convolution to the input tensor.
-        then merge the output tensor to the desired shape.
+        Apply the temporal projection to the input tensor after grouping channels.
 
-        Parameters:
-        -----------
-        x: torch.Tensor
-            Input tensor of shape (Batch, Channels, n_patchs, patch_size).
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (Batch, n_channels, n_times) or
+            (Batch, n_channels, n_patches, patch_size).
 
-        Return:
+        Returns
         -------
-        x: torch.Tensor
-            Output tensor of shape (Batch, n_patchs, patch_size, channels).
+        torch.Tensor
+            Output tensor of shape (Batch, n_patchs, emb_dim).
         """
+        if x.ndim == 4:
+            batch_size, n_channels, n_patchs, patch_len = x.shape
+            if patch_len != self.patch_size:
+                raise ValueError(
+                    "When providing a 4D tensor, the last dimension "
+                    f"({patch_len}) must match patch_size ({self.patch_size})."
+                )
+            n_times = n_patchs * patch_len
+            x = x.reshape(batch_size, n_channels, n_times)
+        elif x.ndim == 3:
+            batch_size, n_channels, n_times = x.shape
+        else:
+            raise ValueError(
+                "Input must be either 3D (batch, channels, times) or "
+                "4D (batch, channels, n_patches, patch_size)."
+            )
+
+        if n_times % self.patch_size != 0:
+            raise ValueError(
+                f"n_times ({n_times}) must be divisible by patch_size ({self.patch_size})."
+            )
+        if n_channels % self.in_channels != 0:
+            raise ValueError(
+                "The input channel dimension "
+                f"({n_channels}) must be divisible by in_channels ({self.in_channels})."
+            )
+
+        group_size = n_channels // self.in_channels
+
+        # Reshape so Conv2d sees `in_channels` feature maps and uses the grouped
+        # EEG channels as the spatial height dimension.
+        # Shape after view: (Batch, in_channels, group_size, n_times)
+        x = x.view(batch_size, self.in_channels, group_size, n_times)
+
+        # Apply the temporal projection per group.
+        # Output shape: (Batch, emb_dim, group_size, n_patchs)
         x = self.proj(x)
-        x = self.merge_transpose(x)
+
+        # THIS IS braindecode's MODIFICATION:
+        # Average over the grouped channel dimension and permute to (Batch, n_patchs, emb_dim)
+        x = x.mean(dim=2)
+        x = x.transpose(1, 2).contiguous()
+
         return x