braindecode 1.3.0.dev177069446__py3-none-any.whl → 1.3.0.dev177628147__py3-none-any.whl

braindecode/models/biot.py

@@ -9,9 +9,9 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class BIOT(EEGModuleMixin, nn.Module):
-    r"""BIOT from Yang et al (2023) [Yang2023]_
+    """BIOT from Yang et al. (2023) [Yang2023]_
 
-    :bdg-danger:`Foundation Model`
+    :bdg-danger:`Large Brain Model`
 
     .. figure:: https://braindecode.org/dev/_static/model/biot.jpg
        :align: center
@@ -19,7 +19,7 @@ class BIOT(EEGModuleMixin, nn.Module):
 
     BIOT: Cross-data Biosignal Learning in the Wild.
 
-    BIOT is a foundation model for biosignal classification. It is
+    BIOT is a large brain model for biosignal classification. It is
     a wrapper around the `BIOTEncoder` and `ClassificationHead` modules.
 
     It is designed for N-dimensional biosignal data such as EEG, ECG, etc.
@@ -41,44 +41,15 @@ class BIOT(EEGModuleMixin, nn.Module):
     linear layer that takes the output of the `BIOTEncoder` and outputs
     the classification probabilities.
 
-    .. important::
-       **Pre-trained Weights Available**
-
-       This model has pre-trained weights available on the Hugging Face Hub.
-       You can load them using:
-
-       .. code-block:: python
-
-           from braindecode.models import BIOT
-
-           # Load the original pre-trained model from Hugging Face Hub
-           # For 16-channel models:
-           model = BIOT.from_pretrained("braindecode/biot-pretrained-prest-16chs")
-
-           # For 18-channel models:
-           model = BIOT.from_pretrained("braindecode/biot-pretrained-shhs-prest-18chs")
-           model = BIOT.from_pretrained("braindecode/biot-pretrained-six-datasets-18chs")
-
-       To push your own trained model to the Hub:
-
-       .. code-block:: python
-
-           # After training your model
-           model.push_to_hub(
-               repo_id="username/my-biot-model", commit_message="Upload trained BIOT model"
-           )
-
-       Requires installing ``braindecode[hug]`` for Hub integration.
-
     .. versionadded:: 0.9
 
     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
@@ -105,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,
@@ -116,12 +87,12 @@ class BIOT(EEGModuleMixin, nn.Module):
         chs_info=None,
         n_times=None,
         input_window_seconds=None,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
         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,
@@ -132,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(
@@ -143,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 "
@@ -171,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,
         )
@@ -216,7 +187,7 @@ class BIOT(EEGModuleMixin, nn.Module):
 
 
 class _PatchFrequencyEmbedding(nn.Module):
-    r"""
+    """
     Patch Frequency Embedding.
 
     A simple linear layer is used to learn some representation over the
@@ -258,7 +229,7 @@ class _PatchFrequencyEmbedding(nn.Module):
 
 
 class _ClassificationHead(nn.Sequential):
-    r"""
+    """
     Classification head for the BIOT model.
 
     Simple linear layer with ELU activation function.
@@ -279,9 +250,7 @@ class _ClassificationHead(nn.Sequential):
         (batch, n_outputs)
     """
 
-    def __init__(
-        self, emb_size: int, n_outputs: int, activation: type[nn.Module] = nn.ELU
-    ):
+    def __init__(self, emb_size: int, n_outputs: int, activation: nn.Module = nn.ELU):
         super().__init__()
         self.activation_layer = activation()
         self.classification_head = nn.Linear(emb_size, n_outputs)
@@ -293,7 +262,7 @@ class _ClassificationHead(nn.Sequential):
 
 
 class _PositionalEncoding(nn.Module):
-    r"""
+    """
     Positional Encoding.
 
     We first create a `pe` zero matrix of shape (max_len, d_model) where max_len is the
@@ -354,7 +323,7 @@ class _PositionalEncoding(nn.Module):
 
 
 class _BIOTEncoder(nn.Module):
-    r"""
+    """
     BIOT Encoder.
 
     The BIOT encoder is a transformer that takes the time series input data and
@@ -376,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
@@ -389,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
@@ -409,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/models/contrawr.py

@@ -8,7 +8,7 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class ContraWR(EEGModuleMixin, nn.Module):
-    r"""Contrast with the World Representation ContraWR from Yang et al (2021) [Yang2021]_.
+    """Contrast with the World Representation ContraWR from Yang et al (2021) [Yang2021]_.
 
     :bdg-success:`Convolution`
 
@@ -58,7 +58,7 @@ class ContraWR(EEGModuleMixin, nn.Module):
         emb_size: int = 256,
         res_channels: list[int] = [32, 64, 128],
         steps=20,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
         drop_prob: float = 0.5,
         stride_res: int = 2,
         kernel_size_res: int = 3,
@@ -148,7 +148,7 @@ class ContraWR(EEGModuleMixin, nn.Module):
 
 
 class _ResBlock(nn.Module):
-    r"""Convolutional Residual Block 2D.
+    """Convolutional Residual Block 2D.
 
     This block stacks two convolutional layers with batch normalization,
     max pooling, dropout, and residual connection.
@@ -195,7 +195,7 @@ class _ResBlock(nn.Module):
         kernel_size=3,
         padding=1,
         drop_prob=0.5,
-        activation: type[nn.Module] = nn.ReLU,
+        activation: nn.Module = nn.ReLU,
     ):
         super().__init__()
         self.conv1 = nn.Conv2d(
@@ -259,7 +259,7 @@ class _ResBlock(nn.Module):
 
 
 class _STFTModule(nn.Module):
-    r"""
+    """
     A PyTorch module that computes the Short-Time Fourier Transform (STFT)
     of an EEG batch tensor.
 

braindecode/models/ctnet.py

@@ -25,9 +25,9 @@ from braindecode.modules import (
 
 
 class CTNet(EEGModuleMixin, nn.Module):
-    r"""CTNet from Zhao, W et al (2024) [ctnet]_.
+    """CTNet from Zhao, W et al (2024) [ctnet]_.
 
-    :bdg-success:`Convolution` :bdg-info:`Attention/Transformer`
+    :bdg-success:`Convolution` :bdg-info:`Small Attention`
 
      A Convolutional Transformer Network for EEG-Based Motor Imagery Classification
 
@@ -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
@@ -109,15 +109,15 @@ class CTNet(EEGModuleMixin, nn.Module):
         n_times=None,
         input_window_seconds=None,
         # Model specific arguments
-        activation_patch: type[nn.Module] = nn.ELU,
-        activation_transformer: type[nn.Module] = nn.GELU,
-        cnn_drop_prob: float = 0.3,
-        att_positional_drop_prob: float = 0.1,
-        final_drop_prob: float = 0.5,
+        activation_patch: nn.Module = nn.ELU,
+        activation_transformer: nn.Module = nn.GELU,
+        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
@@ -312,7 +312,7 @@ class _PatchEmbeddingEEGNet(nn.Module):
         pool_size_2: int = 8,
         drop_prob: float = 0.3,
         n_chans: int = 22,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
     ):
         super().__init__()
         n_filters_out = depth_multiplier * n_filters_time
@@ -416,7 +416,7 @@ class _TransformerEncoderBlock(nn.Module):
         drop_prob: float = 0.5,
         forward_expansion: int = 4,
         forward_drop_p: float = 0.5,
-        activation: type[nn.Module] = nn.GELU,
+        activation: nn.Module = nn.GELU,
     ):
         super().__init__()
         self.attention = _ResidualAdd(
@@ -466,7 +466,7 @@ class _TransformerEncoder(nn.Module):
         nheads: int,
         depth: int,
         dim_feedforward: int,
-        activation: type[nn.Module] = nn.GELU,
+        activation: nn.Module = nn.GELU,
     ):
         super().__init__()
         self.layers = nn.Sequential(

braindecode/models/deep4.py

@@ -17,7 +17,7 @@ from braindecode.modules import (
 
 
 class Deep4Net(EEGModuleMixin, nn.Sequential):
-    r"""Deep ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
+    """Deep ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
 
     :bdg-success:`Convolution`
 
@@ -109,12 +109,12 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         filter_length_3=10,
         n_filters_4=200,
         filter_length_4=10,
-        activation_first_conv_nonlin: type[nn.Module] = nn.ELU,
+        activation_first_conv_nonlin: nn.Module = nn.ELU,
         first_pool_mode="max",
-        first_pool_nonlin: type[nn.Module] = nn.Identity,
-        activation_later_conv_nonlin: type[nn.Module] = nn.ELU,
+        first_pool_nonlin: nn.Module = nn.Identity,
+        activation_later_conv_nonlin: nn.Module = nn.ELU,
         later_pool_mode="max",
-        later_pool_nonlin: type[nn.Module] = nn.Identity,
+        later_pool_nonlin: nn.Module = nn.Identity,
         drop_prob=0.5,
         split_first_layer=True,
         batch_norm=True,

braindecode/models/deepsleepnet.py

@@ -8,7 +8,7 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class DeepSleepNet(EEGModuleMixin, nn.Module):
-    r"""DeepSleepNet from Supratak et al (2017) [Supratak2017]_.
+    """DeepSleepNet from Supratak et al. (2017) [Supratak2017]_.
 
     :bdg-success:`Convolution` :bdg-secondary:`Recurrent`
 
@@ -172,8 +172,8 @@ class DeepSleepNet(EEGModuleMixin, nn.Module):
         n_times=None,
         input_window_seconds=None,
         sfreq=None,
-        activation_large: type[nn.Module] = nn.ELU,
-        activation_small: type[nn.Module] = nn.ReLU,
+        activation_large: nn.Module = nn.ELU,
+        activation_small: nn.Module = nn.ReLU,
         drop_prob: float = 0.5,
     ):
         super().__init__(
@@ -240,7 +240,7 @@ class DeepSleepNet(EEGModuleMixin, nn.Module):
 
 
 class _SmallCNN(nn.Module):
-    r"""
+    """
     Smaller filter sizes to learn temporal information.
 
     Parameters
@@ -252,7 +252,7 @@ class _SmallCNN(nn.Module):
         The dropout rate for regularization. Values should be between 0 and 1.
     """
 
-    def __init__(self, activation: type[nn.Module] = nn.ReLU, drop_prob: float = 0.5):
+    def __init__(self, activation: nn.Module = nn.ReLU, drop_prob: float = 0.5):
         super().__init__()
         self.conv1 = nn.Sequential(
             nn.Conv2d(
@@ -317,7 +317,7 @@ class _SmallCNN(nn.Module):
 
 
 class _LargeCNN(nn.Module):
-    r"""
+    """
     Larger filter sizes to learn frequency information.
 
     Parameters
@@ -328,7 +328,7 @@ class _LargeCNN(nn.Module):
 
     """
 
-    def __init__(self, activation: type[nn.Module] = nn.ELU, drop_prob: float = 0.5):
+    def __init__(self, activation: nn.Module = nn.ELU, drop_prob: float = 0.5):
         super().__init__()
 
         self.conv1 = nn.Sequential(

braindecode/models/eegconformer.py

@@ -12,9 +12,9 @@ from braindecode.modules import FeedForwardBlock, MultiHeadAttention
 
 
 class EEGConformer(EEGModuleMixin, nn.Module):
-    r"""EEG Conformer from Song et al (2022) [song2022]_.
+    """EEG Conformer from Song et al. (2022) [song2022]_.
 
-    :bdg-success:`Convolution` :bdg-info:`Attention/Transformer`
+    :bdg-success:`Convolution` :bdg-info:`Small Attention`
 
     .. figure:: https://raw.githubusercontent.com/eeyhsong/EEG-Conformer/refs/heads/main/visualization/Fig1.png
         :align: center
@@ -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,13 +197,13 @@ 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,
-        activation: type[nn.Module] = nn.ELU,
-        activation_transfor: type[nn.Module] = nn.GELU,
+        activation: nn.Module = nn.ELU,
+        activation_transfor: nn.Module = nn.GELU,
         n_times=None,
         chs_info=None,
         input_window_seconds=None,
@@ -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,
         )
@@ -284,7 +284,7 @@ class EEGConformer(EEGModuleMixin, nn.Module):
 
 
 class _PatchEmbedding(nn.Module):
-    r"""Patch Embedding.
+    """Patch Embedding.
 
     The authors used a convolution module to capture local features,
     instead of position embedding.
@@ -318,7 +318,7 @@ class _PatchEmbedding(nn.Module):
         pool_time_length,
         stride_avg_pool,
         drop_prob,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
     ):
         super().__init__()
 
@@ -364,16 +364,16 @@ class _TransformerEncoderBlock(nn.Sequential):
     def __init__(
         self,
         emb_size,
-        num_heads,
+        att_heads,
         att_drop,
         forward_expansion=4,
-        activation: type[nn.Module] = nn.GELU,
+        activation: nn.Module = nn.GELU,
     ):
         super().__init__(
             _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),
                 )
             ),
@@ -393,17 +393,17 @@ class _TransformerEncoderBlock(nn.Sequential):
 
 
 class _TransformerEncoder(nn.Sequential):
-    r"""Transformer encoder module for the transformer encoder.
+    """Transformer encoder module for the transformer encoder.
 
     Similar to the layers used in ViT.
 
     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,19 +411,14 @@ class _TransformerEncoder(nn.Sequential):
     """
 
     def __init__(
-        self,
-        num_layers,
-        emb_size,
-        num_heads,
-        att_drop,
-        activation: type[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)
             ]
         )
 
@@ -436,7 +431,7 @@ class _FullyConnected(nn.Module):
         drop_prob_2=0.3,
         out_channels=256,
         hidden_channels=32,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
     ):
         """Fully-connected layer for the transformer encoder.