Braindecode 1.3.0.dev169126795__tar.gz → 1.3.0.dev169361694__tar.gz

{braindecode-1.3.0.dev169126795/braindecode.egg-info → braindecode-1.3.0.dev169361694}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.0.dev169126795
+Version: 1.3.0.dev169361694
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>

{braindecode-1.3.0.dev169126795 → braindecode-1.3.0.dev169361694}/braindecode/models/atcnet.py

@@ -141,7 +141,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
     - Pool sizes ``P1,P2`` trade temporal resolution for stability/compute; they set
       ``T_c = T/(P1·P2)`` and thus window width ``T_w``.
     - ``n_windows`` controls the ensemble over shifts (compute ∝ windows).
-    - ``num_heads``, ``head_dim`` set attention capacity; keep ``H·d_h ≈ F2``.
+    - ``att_num_heads``, ``att_head_dim`` set attention capacity; keep ``H·d_h ≈ F2``.
     - ``tcn_depth``, ``tcn_kernel_size`` govern receptive field; larger values demand
       longer inputs (see minimum length above). The implementation warns and *rescales*
       kernels/pools/windows if inputs are too short.
@@ -194,10 +194,10 @@ class ATCNet(EEGModuleMixin, nn.Module):
         table 1 of the paper [1]_. Defaults to 0.3 as in [1]_.
     n_windows : int
         Number of sliding windows, denoted n in [1]_. Defaults to 5 as in [1]_.
-    head_dim : int
+    att_head_dim : int
         Embedding dimension used in each self-attention head, denoted dh in
         table 1 of the paper [1]_. Defaults to 8 as in [1]_.
-    num_heads : int
+    att_num_heads : int
         Number of attention heads, denoted H in table 1 of the paper [1]_.
         Defaults to 2 as in [1]_.
     att_dropout : float
@@ -248,8 +248,8 @@ class ATCNet(EEGModuleMixin, nn.Module):
         conv_block_depth_mult=2,
         conv_block_dropout=0.3,
         n_windows=5,
-        head_dim=8,
-        num_heads=2,
+        att_head_dim=8,
+        att_num_heads=2,
         att_drop_prob=0.5,
         tcn_depth=2,
         tcn_kernel_size=4,
@@ -316,8 +316,8 @@ class ATCNet(EEGModuleMixin, nn.Module):
         self.conv_block_depth_mult = conv_block_depth_mult
         self.conv_block_dropout = conv_block_dropout
         self.n_windows = n_windows
-        self.head_dim = head_dim
-        self.num_heads = num_heads
+        self.att_head_dim = att_head_dim
+        self.att_num_heads = att_num_heads
         self.att_dropout = att_drop_prob
         self.tcn_depth = tcn_depth
         self.tcn_kernel_size = tcn_kernel_size
@@ -356,8 +356,8 @@ class ATCNet(EEGModuleMixin, nn.Module):
             [
                 _AttentionBlock(
                     in_shape=self.F2,
-                    head_dim=self.head_dim,
-                    num_heads=num_heads,
+                    head_dim=self.att_head_dim,
+                    num_heads=att_num_heads,
                     dropout=att_drop_prob,
                 )
                 for _ in range(self.n_windows)

{braindecode-1.3.0.dev169126795 → braindecode-1.3.0.dev169361694}/braindecode/models/biot.py

@@ -45,11 +45,11 @@ class BIOT(EEGModuleMixin, nn.Module):
 
     Parameters
     ----------
-    embed_dim : int, optional
+    emb_size : int, optional
         The size of the embedding layer, by default 256
-    num_heads : int, optional
+    att_num_heads : int, optional
         The number of attention heads, by default 8
-    num_layers : int, optional
+    n_layers : int, optional
         The number of transformer layers, by default 4
     activation: nn.Module, default=nn.ELU
         Activation function class to apply. Should be a PyTorch activation
@@ -76,9 +76,9 @@ class BIOT(EEGModuleMixin, nn.Module):
 
     def __init__(
         self,
-        embed_dim=256,
-        num_heads=8,
-        num_layers=4,
+        emb_size=256,
+        att_num_heads=8,
+        n_layers=4,
         sfreq=200,
         hop_length=100,
         return_feature=False,
@@ -91,8 +91,8 @@ class BIOT(EEGModuleMixin, nn.Module):
         drop_prob: float = 0.5,
         # Parameters for the encoder
         max_seq_len: int = 1024,
-        att_drop_prob=0.2,
-        att_layer_drop_prob=0.2,
+        attn_dropout=0.2,
+        attn_layer_dropout=0.2,
     ):
         super().__init__(
             n_outputs=n_outputs,
@@ -103,10 +103,10 @@ class BIOT(EEGModuleMixin, nn.Module):
             sfreq=sfreq,
         )
         del n_outputs, n_chans, chs_info, n_times, sfreq
-        self.embed_dim = embed_dim
+        self.emb_size = emb_size
         self.hop_length = hop_length
-        self.num_heads = num_heads
-        self.num_layers = num_layers
+        self.att_num_heads = att_num_heads
+        self.n_layers = n_layers
         self.return_feature = return_feature
         if (self.sfreq != 200) & (self.sfreq is not None):
             warn(
@@ -114,7 +114,7 @@ class BIOT(EEGModuleMixin, nn.Module):
                 + "no guarantee to generalize well with the default parameters",
                 UserWarning,
             )
-        if self.n_chans > embed_dim:
+        if self.n_chans > emb_size:
             warn(
                 "The number of channels is larger than the embedding size. "
                 + "This may cause overfitting. Consider using a larger "
@@ -142,20 +142,20 @@ class BIOT(EEGModuleMixin, nn.Module):
             self.n_fft = int(self.sfreq)
 
         self.encoder = _BIOTEncoder(
-            emb_size=self.embed_dim,
-            num_heads=self.num_heads,
-            n_layers=self.num_layers,
+            emb_size=emb_size,
+            att_num_heads=att_num_heads,
+            n_layers=n_layers,
             n_chans=self.n_chans,
             n_fft=self.n_fft,
             hop_length=hop_length,
             drop_prob=drop_prob,
             max_seq_len=max_seq_len,
-            attn_dropout=att_drop_prob,
-            attn_layer_dropout=att_layer_drop_prob,
+            attn_dropout=attn_dropout,
+            attn_layer_dropout=attn_layer_dropout,
         )
 
         self.final_layer = _ClassificationHead(
-            emb_size=self.embed_dim,
+            emb_size=emb_size,
             n_outputs=self.n_outputs,
             activation=activation,
         )
@@ -345,7 +345,7 @@ class _BIOTEncoder(nn.Module):
         The number of channels
     emb_size: int
         The size of the embedding layer
-    num_heads: int
+    att_num_heads: int
         The number of attention heads
     n_layers: int
         The number of transformer layers
@@ -358,7 +358,7 @@ class _BIOTEncoder(nn.Module):
     def __init__(
         self,
         emb_size=256,  # The size of the embedding layer
-        num_heads=8,  # The number of attention heads
+        att_num_heads=8,  # The number of attention heads
         n_chans=16,  # The number of channels
         n_layers=4,  # The number of transformer layers
         n_fft=200,  # Related with the frequency resolution
@@ -378,7 +378,7 @@ class _BIOTEncoder(nn.Module):
         )
         self.transformer = LinearAttentionTransformer(
             dim=emb_size,
-            heads=num_heads,
+            heads=att_num_heads,
             depth=n_layers,
             max_seq_len=max_seq_len,
             attn_layer_dropout=attn_layer_dropout,

{braindecode-1.3.0.dev169126795 → braindecode-1.3.0.dev169361694}/braindecode/models/ctnet.py

@@ -61,11 +61,11 @@ class CTNet(EEGModuleMixin, nn.Module):
     ----------
     activation : nn.Module, default=nn.GELU
         Activation function to use in the network.
-    num_heads : int, default=4
+    heads : int, default=4
         Number of attention heads in the Transformer encoder.
-    embed_dim : int or None, default=None
+    emb_size : int or None, default=None
         Embedding size (dimensionality) for the Transformer encoder.
-    num_layers : int, default=6
+    depth : int, default=6
         Number of encoder layers in the Transformer.
     n_filters_time : int, default=20
         Number of temporal filters in the first convolutional layer.
@@ -77,11 +77,11 @@ class CTNet(EEGModuleMixin, nn.Module):
         Pooling size for the first average pooling layer.
     pool_size_2 : int, default=8
         Pooling size for the second average pooling layer.
-        cnn_drop_prob: float, default=0.3
+    drop_prob_cnn : float, default=0.3
         Dropout probability after convolutional layers.
-    att_positional_drop_prob : float, default=0.1
+    drop_prob_posi : float, default=0.1
         Dropout probability for the positional encoding in the Transformer.
-    final_drop_prob : float, default=0.5
+    drop_prob_final : float, default=0.5
         Dropout probability before the final classification layer.
 
     Notes
@@ -111,13 +111,13 @@ class CTNet(EEGModuleMixin, nn.Module):
         # Model specific arguments
         activation_patch: nn.Module = nn.ELU,
         activation_transformer: nn.Module = nn.GELU,
-        cnn_drop_prob: float = 0.3,
-        att_positional_drop_prob: float = 0.1,
-        final_drop_prob: float = 0.5,
+        drop_prob_cnn: float = 0.3,
+        drop_prob_posi: float = 0.1,
+        drop_prob_final: float = 0.5,
         # other parameters
-        num_heads: int = 4,
-        embed_dim: Optional[int] = 40,
-        num_layers: int = 6,
+        heads: int = 4,
+        emb_size: Optional[int] = 40,
+        depth: int = 6,
         n_filters_time: Optional[int] = None,
         kernel_size: int = 64,
         depth_multiplier: Optional[int] = 2,
@@ -136,14 +136,14 @@ class CTNet(EEGModuleMixin, nn.Module):
 
         self.activation_patch = activation_patch
         self.activation_transformer = activation_transformer
-        self.cnn_drop_prob = cnn_drop_prob
+        self.drop_prob_cnn = drop_prob_cnn
         self.pool_size_1 = pool_size_1
         self.pool_size_2 = pool_size_2
         self.kernel_size = kernel_size
-        self.att_positional_drop_prob = att_positional_drop_prob
-        self.final_drop_prob = final_drop_prob
-        self.num_heads = num_heads
-        self.num_layers = num_layers
+        self.drop_prob_posi = drop_prob_posi
+        self.drop_prob_final = drop_prob_final
+        self.heads = heads
+        self.depth = depth
         # n_times - pool_size_1 / p
         self.sequence_length = math.floor(
             (
@@ -154,8 +154,8 @@ class CTNet(EEGModuleMixin, nn.Module):
             + 1
         )
 
-        self.depth_multiplier, self.n_filters_time, self.embed_dim = self._resolve_dims(
-            depth_multiplier, n_filters_time, embed_dim
+        self.depth_multiplier, self.n_filters_time, self.emb_size = self._resolve_dims(
+            depth_multiplier, n_filters_time, emb_size
         )
 
         # Layers
@@ -168,32 +168,32 @@ class CTNet(EEGModuleMixin, nn.Module):
             depth_multiplier=self.depth_multiplier,
             pool_size_1=self.pool_size_1,
             pool_size_2=self.pool_size_2,
-            drop_prob=self.cnn_drop_prob,
+            drop_prob=self.drop_prob_cnn,
             n_chans=self.n_chans,
             activation=self.activation_patch,
         )
 
         self.position = _PositionalEncoding(
-            emb_size=self.embed_dim,
-            drop_prob=self.att_positional_drop_prob,
+            emb_size=self.emb_size,
+            drop_prob=self.drop_prob_posi,
             n_times=self.n_times,
             pool_size=self.pool_size_1,
         )
 
         self.trans = _TransformerEncoder(
-            self.num_heads,
-            self.num_layers,
-            self.embed_dim,
+            self.heads,
+            self.depth,
+            self.emb_size,
             activation=self.activation_transformer,
         )
 
         self.flatten_drop_layer = nn.Sequential(
             nn.Flatten(),
-            nn.Dropout(p=self.final_drop_prob),
+            nn.Dropout(p=self.drop_prob_final),
         )
 
         self.final_layer = nn.Linear(
-            in_features=int(self.embed_dim * self.sequence_length),
+            in_features=int(self.emb_size * self.sequence_length),
             out_features=self.n_outputs,
         )
 
@@ -213,7 +213,7 @@ class CTNet(EEGModuleMixin, nn.Module):
         """
         x = self.ensuredim(x)
         cnn = self.cnn(x)
-        cnn = cnn * math.sqrt(self.embed_dim)
+        cnn = cnn * math.sqrt(self.emb_size)
         cnn = self.position(cnn)
         trans = self.trans(cnn)
         features = cnn + trans

{braindecode-1.3.0.dev169126795 → braindecode-1.3.0.dev169361694}/braindecode/models/eegconformer.py

@@ -57,9 +57,9 @@ class EEGConformer(EEGModuleMixin, nn.Module):
     - :class:`_TransformerEncoder` **(context over temporal tokens)**
 
         - *Operations.*
-        - A stack of ``num_layers`` encoder blocks. :class:`_TransformerEncoderBlock`
+        - A stack of ``att_depth`` encoder blocks. :class:`_TransformerEncoderBlock`
         - Each block applies LayerNorm :class:`torch.nn.LayerNorm`
-        - Multi-Head Self-Attention (``num_heads``) with dropout + residual :class:`MultiHeadAttention` (:class:`torch.nn.Dropout`)
+        - Multi-Head Self-Attention (``att_heads``) with dropout + residual :class:`MultiHeadAttention` (:class:`torch.nn.Dropout`)
         - LayerNorm :class:`torch.nn.LayerNorm`
         - 2-layer feed-forward (≈4x expansion, :class:`torch.nn.GELU`) with dropout + residual.
 
@@ -100,7 +100,7 @@ class EEGConformer(EEGModuleMixin, nn.Module):
 
     .. rubric:: Attention / Sequential Modules
 
-    - **Type.** Standard multi-head self-attention (MHA) with ``num_heads`` heads over the token sequence.
+    - **Type.** Standard multi-head self-attention (MHA) with ``att_heads`` heads over the token sequence.
     - **Shapes.** Input/Output: ``(B, S_tokens, D)``; attention operates along the ``S_tokens`` axis.
     - **Role.** Re-weights and integrates evidence across pooled windows, capturing dependencies
       longer than any single token while leaving channel relationships to the convolutional stem.
@@ -127,7 +127,7 @@ class EEGConformer(EEGModuleMixin, nn.Module):
     - **Instantiation.** Choose ``n_filters_time`` (embedding size ``D``) and
         ``filter_time_length`` to match the rhythms of interest. Tune
         ``pool_time_length/stride`` to trade temporal resolution for sequence length.
-        Keep ``num_layers`` modest (e.g., 4–6) and set ``num_heads`` to divide ``D``.
+        Keep ``att_depth`` modest (e.g., 4–6) and set ``att_heads`` to divide ``D``.
         ``final_fc_length="auto"`` infers the flattened size from PatchEmbedding.
 
     Notes
@@ -160,9 +160,9 @@ class EEGConformer(EEGModuleMixin, nn.Module):
         Length of stride between temporal pooling filters.
     drop_prob: float
         Dropout rate of the convolutional layer.
-    num_layers: int
+    att_depth: int
         Number of self-attention layers.
-    num_heads: int
+    att_heads: int
         Number of attention heads.
     att_drop_prob: float
         Dropout rate of the self-attention layer.
@@ -197,8 +197,8 @@ class EEGConformer(EEGModuleMixin, nn.Module):
         pool_time_length=75,
         pool_time_stride=15,
         drop_prob=0.5,
-        num_layers=6,
-        num_heads=10,
+        att_depth=6,
+        att_heads=10,
         att_drop_prob=0.5,
         final_fc_length="auto",
         return_features=False,
@@ -250,9 +250,9 @@ class EEGConformer(EEGModuleMixin, nn.Module):
             self.final_fc_length = final_fc_length
 
         self.transformer = _TransformerEncoder(
-            num_layers=num_layers,
+            att_depth=att_depth,
             emb_size=n_filters_time,
-            num_heads=num_heads,
+            att_heads=att_heads,
             att_drop=att_drop_prob,
             activation=activation_transfor,
         )
@@ -364,7 +364,7 @@ class _TransformerEncoderBlock(nn.Sequential):
     def __init__(
         self,
         emb_size,
-        num_heads,
+        att_heads,
         att_drop,
         forward_expansion=4,
         activation: nn.Module = nn.GELU,
@@ -373,7 +373,7 @@ class _TransformerEncoderBlock(nn.Sequential):
             _ResidualAdd(
                 nn.Sequential(
                     nn.LayerNorm(emb_size),
-                    MultiHeadAttention(emb_size, num_heads, att_drop),
+                    MultiHeadAttention(emb_size, att_heads, att_drop),
                     nn.Dropout(att_drop),
                 )
             ),
@@ -399,11 +399,11 @@ class _TransformerEncoder(nn.Sequential):
 
     Parameters
     ----------
-    num_layers : int
+    att_depth : int
         Number of transformer encoder blocks.
     emb_size : int
         Embedding size of the transformer encoder.
-    num_heads : int
+    att_heads : int
         Number of attention heads.
     att_drop : float
         Dropout probability for the attention layers.
@@ -411,14 +411,14 @@ class _TransformerEncoder(nn.Sequential):
     """
 
     def __init__(
-        self, num_layers, emb_size, num_heads, att_drop, activation: nn.Module = nn.GELU
+        self, att_depth, emb_size, att_heads, att_drop, activation: nn.Module = nn.GELU
     ):
         super().__init__(
             *[
                 _TransformerEncoderBlock(
-                    emb_size, num_heads, att_drop, activation=activation
+                    emb_size, att_heads, att_drop, activation=activation
                 )
-                for _ in range(num_layers)
+                for _ in range(att_depth)
             ]
         )
 

{braindecode-1.3.0.dev169126795 → braindecode-1.3.0.dev169361694}/braindecode/models/labram.py

@@ -83,15 +83,15 @@ class Labram(EEGModuleMixin, nn.Module):
     ----------
     patch_size : int
         The size of the patch to be used in the patch embedding.
-    embed_dim : int
+    emb_size : int
         The dimension of the embedding.
-    conv_in_channels : int
+    in_conv_channels : int
         The number of convolutional input channels.
-    conv_out_channels : int
+    out_channels : int
         The number of convolutional output channels.
-    num_layers :  int (default=12)
+    n_layers :  int (default=12)
         The number of attention layers of the model.
-    num_heads : int (default=10)
+    att_num_heads : int (default=10)
         The number of attention heads.
     mlp_ratio : float (default=4.0)
         The expansion ratio of the mlp layer
@@ -155,11 +155,11 @@ class Labram(EEGModuleMixin, nn.Module):
         sfreq=None,
         input_window_seconds=None,
         patch_size=200,
-        embed_dim=200,
-        conv_in_channels=1,
-        conv_out_channels=8,
-        num_layers=12,
-        num_heads=10,
+        emb_size=200,
+        in_conv_channels=1,
+        out_channels=8,
+        n_layers=12,
+        att_num_heads=10,
         mlp_ratio=4.0,
         qkv_bias=False,
         qk_norm=nn.LayerNorm,
@@ -187,7 +187,7 @@ class Labram(EEGModuleMixin, nn.Module):
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
 
         self.patch_size = patch_size
-        self.num_features = self.embed_dim = embed_dim
+        self.num_features = self.emb_size = emb_size
         self.neural_tokenizer = neural_tokenizer
         self.init_scale = init_scale
 
@@ -199,20 +199,20 @@ class Labram(EEGModuleMixin, nn.Module):
             )
             self.patch_size = self.n_times
             self.num_features = None
-            self.embed_dim = None
+            self.emb_size = None
         else:
             self.patch_size = patch_size
         self.n_path = self.n_times // self.patch_size
 
-        if neural_tokenizer and conv_in_channels != 1:
+        if neural_tokenizer and in_conv_channels != 1:
             warn(
                 "The model is in Neural Tokenizer mode, but the variable "
-                + "`conv_in_channels` is different from the default values."
-                + "`conv_in_channels` is only needed for the Neural Decoder mode."
-                + "conv_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,
             )
-            conv_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:
@@ -231,7 +231,7 @@ class Labram(EEGModuleMixin, nn.Module):
                         (
                             "temporal_conv",
                             _TemporalConv(
-                                out_channels=conv_out_channels, activation=activation
+                                out_channels=out_channels, activation=activation
                             ),
                         ),
                     ]
@@ -249,8 +249,8 @@ class Labram(EEGModuleMixin, nn.Module):
                 _PatchEmbed(
                     n_times=self.n_times,
                     patch_size=patch_size,
-                    in_channels=conv_in_channels,
-                    emb_dim=self.embed_dim,
+                    in_channels=in_conv_channels,
+                    emb_dim=self.emb_size,
                 ),
             )
 
@@ -259,12 +259,12 @@ class Labram(EEGModuleMixin, nn.Module):
             out = self.patch_embed(dummy)
         # out.shape for tokenizer: (1, n_chans, emb_dim)
         # for decoder:        (1, n_patch, patch_size, emb_dim), but we want last dim
-        self.embed_dim = out.shape[-1]
-        self.num_features = self.embed_dim
+        self.emb_size = out.shape[-1]
+        self.num_features = self.emb_size
 
         # Defining the parameters
         # Creating a parameter list with cls token]
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.emb_size))
         # Positional embedding and time embedding are complementary
         # one is for the spatial information and the other is for the temporal
         # information.
@@ -273,26 +273,26 @@ class Labram(EEGModuleMixin, nn.Module):
         # information.
         if use_abs_pos_emb:
             self.position_embedding = nn.Parameter(
-                torch.zeros(1, self.n_chans + 1, self.embed_dim),
+                torch.zeros(1, self.n_chans + 1, self.emb_size),
                 requires_grad=True,
             )
         else:
             self.position_embedding = None
 
         self.temporal_embedding = nn.Parameter(
-            torch.zeros(1, self.patch_embed[0].n_patchs + 1, self.embed_dim),
+            torch.zeros(1, self.patch_embed[0].n_patchs + 1, self.emb_size),
             requires_grad=True,
         )
         self.pos_drop = nn.Dropout(p=drop_prob)
 
         dpr = [
-            x.item() for x in torch.linspace(0, drop_path_prob, num_layers)
+            x.item() for x in torch.linspace(0, drop_path_prob, n_layers)
         ]  # stochastic depth decay rule
         self.blocks = nn.ModuleList(
             [
                 _WindowsAttentionBlock(
-                    dim=self.embed_dim,
-                    num_heads=num_heads,
+                    dim=self.emb_size,
+                    num_heads=att_num_heads,
                     mlp_ratio=mlp_ratio,
                     qkv_bias=qkv_bias,
                     qk_norm=qk_norm,
@@ -310,14 +310,14 @@ class Labram(EEGModuleMixin, nn.Module):
                     attn_head_dim=attn_head_dim,
                     activation=activation,
                 )
-                for i in range(num_layers)
+                for i in range(n_layers)
             ]
         )
-        self.norm = nn.Identity() if use_mean_pooling else norm_layer(self.embed_dim)
-        self.fc_norm = norm_layer(self.embed_dim) if use_mean_pooling else None
+        self.norm = nn.Identity() if use_mean_pooling else norm_layer(self.emb_size)
+        self.fc_norm = norm_layer(self.emb_size) if use_mean_pooling else None
 
         if self.n_outputs > 0:
-            self.final_layer = nn.Linear(self.embed_dim, self.n_outputs)
+            self.final_layer = nn.Linear(self.emb_size, self.n_outputs)
         else:
             self.final_layer = nn.Identity()
 
@@ -415,7 +415,7 @@ class Labram(EEGModuleMixin, nn.Module):
             x = self.patch_embed(x)
             # x shape: (batch, n_chans, emb_dim)
             n_patch = self.n_chans
-            temporal = self.embed_dim
+            temporal = self.emb_size
         else:
             # For neural decoder: input is (batch, n_chans, n_times)
             # patch_embed returns (batch, n_patchs, emb_dim)
@@ -462,7 +462,7 @@ class Labram(EEGModuleMixin, nn.Module):
             # 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.embed_dim
+                temporal = self.emb_size
 
             # Get temporal embeddings for n_patch patches
             n_time_tokens = min(n_patch, self.temporal_embedding.shape[1] - 1)