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

braindecode/models/patchedtransformer.py

@@ -14,9 +14,9 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class PBT(EEGModuleMixin, nn.Module):
-    r"""Patched Brain Transformer (PBT) model from Klein et al (2025) [pbt]_.
+    r"""Patched Brain Transformer (PBT) model from Klein et al. (2025) [pbt]_.
 
-    :bdg-danger:`Foundation Model`
+    :bdg-danger:`Large Brain Model`
 
     This implementation was based in https://github.com/timonkl/PatchedBrainTransformer/
 
@@ -136,7 +136,7 @@ class PBT(EEGModuleMixin, nn.Module):
 
     .. warning::
 
-        **Important:** As the other Foundation Models in Braindecode, :class:`PBT` is
+        **Important:** As the other Large Brain Models in Braindecode, :class:`PBT` is
         designed for large-scale pre-training and fine-tuning. Training from
         scratch on small datasets may lead to suboptimal results. Cross-Dataset
         pre-training and subsequent fine-tuning is recommended to leverage the
@@ -146,9 +146,9 @@ class PBT(EEGModuleMixin, nn.Module):
     ----------
     d_input : int, optional
         Size (in samples) of each patch (token) extracted along the time axis.
-    embed_dim : int, optional
+    d_model : int, optional
         Transformer embedding dimensionality.
-    num_layers : int, optional
+    n_blocks : int, optional
         Number of Transformer encoder layers.
     num_heads : int, optional
         Number of attention heads.
@@ -190,13 +190,13 @@ class PBT(EEGModuleMixin, nn.Module):
         sfreq=None,
         # Model parameters
         d_input: int = 64,
-        embed_dim: int = 128,
-        num_layers: int = 4,
+        d_model: int = 128,
+        n_blocks: int = 4,
         num_heads: int = 4,
         drop_prob: float = 0.1,
         learnable_cls=True,
         bias_transformer=False,
-        activation: type[nn.Module] = nn.GELU,
+        activation: nn.Module = nn.GELU,
     ) -> None:
         super().__init__(
             n_outputs=n_outputs,
@@ -209,8 +209,8 @@ class PBT(EEGModuleMixin, nn.Module):
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
         # Store hyperparameters
         self.d_input = d_input
-        self.embed_dim = embed_dim
-        self.num_layers = num_layers
+        self.d_model = d_model
+        self.n_blocks = n_blocks
         self.num_heads = num_heads
         self.drop_prob = drop_prob
 
@@ -219,11 +219,11 @@ class PBT(EEGModuleMixin, nn.Module):
 
         # Classification token (learnable or fixed zero)
         if learnable_cls:
-            self.cls_token = nn.Parameter(torch.randn(1, 1, self.embed_dim) * 0.002)
+            self.cls_token = nn.Parameter(torch.randn(1, 1, self.d_model) * 0.002)
         else:
             # non-learnable zeroed tensor
             self.cls_token = torch.full(
-                size=(1, 1, self.embed_dim),
+                size=(1, 1, self.d_model),
                 fill_value=0,
                 requires_grad=False,
                 dtype=torch.float32,
@@ -234,20 +234,20 @@ class PBT(EEGModuleMixin, nn.Module):
             n_chans=self.n_chans, n_times=self.n_times, d_input=self.d_input
         )
 
-        # Linear patch projection from token raw-size -> embed_dim
+        # Linear patch projection from token raw-size -> d_model
         self.patching_projection = nn.Linear(
-            in_features=self.d_input, out_features=self.embed_dim, bias=False
+            in_features=self.d_input, out_features=self.d_model, bias=False
         )
 
-        # actual embedding table mapping indices -> embed_dim
+        # actual embedding table mapping indices -> d_model
         self.pos_embedding = nn.Embedding(
-            num_embeddings=self.num_embeddings + 1, embedding_dim=self.embed_dim
+            num_embeddings=self.num_embeddings + 1, embedding_dim=self.d_model
         )
 
         # Transformer encoder stack
         self.transformer_encoder = _TransformerEncoder(
-            num_layers=num_layers,
-            embed_dim=self.embed_dim,
+            n_blocks=n_blocks,
+            d_model=self.d_model,
             n_head=num_heads,
             drop_prob=drop_prob,
             bias=bias_transformer,
@@ -256,7 +256,7 @@ class PBT(EEGModuleMixin, nn.Module):
 
         # classification head on classify token - CLS token
         self.final_layer = nn.Linear(
-            in_features=embed_dim, out_features=self.n_outputs, bias=True
+            in_features=d_model, out_features=self.n_outputs, bias=True
         )
 
         # initialize weights
@@ -305,7 +305,7 @@ class PBT(EEGModuleMixin, nn.Module):
 
 
 class _LayerNorm(nn.Module):
-    r"""Layer normalization with optional bias.
+    """Layer normalization with optional bias.
 
     Simple wrapper around :func:`torch.nn.functional.layer_norm` exposing a
     learnable scale and optional bias.
@@ -346,7 +346,7 @@ class _LayerNorm(nn.Module):
 
 
 class _MHSA(nn.Module):
-    r"""Multi-head self-attention (MHSA) block.
+    """Multi-head self-attention (MHSA) block.
 
     Implements a standard multi-head attention mechanism with optional
     use of PyTorch's scaled_dot_product_attention (FlashAttention) when
@@ -354,7 +354,7 @@ class _MHSA(nn.Module):
 
     Parameters
     ----------
-    embed_dim : int
+    d_model : int
         Dimensionality of the model / embeddings.
     n_head : int
         Number of attention heads.
@@ -366,25 +366,25 @@ class _MHSA(nn.Module):
 
     def __init__(
         self,
-        embed_dim: int,
+        d_model: int,
         n_head: int,
         bias: bool,
         drop_prob: float = 0.0,
     ) -> None:
         super().__init__()
 
-        assert embed_dim % n_head == 0, "embed_dim must be divisible by n_head"
+        assert d_model % n_head == 0, "d_model must be divisible by n_head"
 
         # qkv and output projection
-        self.attn = nn.Linear(embed_dim, 3 * embed_dim, bias=bias)
-        self.proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.attn = nn.Linear(d_model, 3 * d_model, bias=bias)
+        self.proj = nn.Linear(d_model, d_model, bias=bias)
 
         # dropout modules
         self.attn_drop_prob = nn.Dropout(drop_prob)
         self.resid_drop_prob = nn.Dropout(drop_prob)
 
         self.n_head = n_head
-        self.embed_dim = embed_dim
+        self.d_model = d_model
         self.drop_prob = drop_prob
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -393,7 +393,7 @@ class _MHSA(nn.Module):
         Parameters
         ----------
         x : torch.Tensor
-            Input tensor of shape (B, T, C) where C == embed_dim.
+            Input tensor of shape (B, T, C) where C == d_model.
 
         Returns
         -------
@@ -404,7 +404,7 @@ class _MHSA(nn.Module):
         B, T, C = x.size()
 
         # project to q, k, v and reshape for multi-head attention
-        q, k, v = self.attn(x).split(self.embed_dim, dim=2)
+        q, k, v = self.attn(x).split(self.d_model, dim=2)
 
         # (B, nh, T, hs)
         q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
@@ -429,14 +429,14 @@ class _MHSA(nn.Module):
 
 
 class _FeedForward(nn.Module):
-    r"""Position-wise feed-forward network from Transformer.
+    """Position-wise feed-forward network from Transformer.
 
     Implements the two-layer MLP with GELU activation and dropout used in
     Transformer architectures.
 
     Parameters
     ----------
-    embed_dim : int
+    d_model : int
         Input and output dimensionality.
     dim_feedforward : int, optional
         Hidden dimensionality of the feed-forward layer. If None, must be provided by caller.
@@ -448,20 +448,20 @@ class _FeedForward(nn.Module):
 
     def __init__(
         self,
-        embed_dim: int,
+        d_model: int,
         dim_feedforward: Optional[int] = None,
         drop_prob: float = 0.0,
         bias: bool = False,
-        activation: type[nn.Module] = nn.GELU,
+        activation: nn.Module = nn.GELU,
     ) -> None:
         super().__init__()
 
         if dim_feedforward is None:
             raise ValueError("dim_feedforward must be provided")
 
-        self.proj_in = nn.Linear(embed_dim, dim_feedforward, bias=bias)
+        self.proj_in = nn.Linear(d_model, dim_feedforward, bias=bias)
         self.activation = activation()
-        self.proj = nn.Linear(dim_feedforward, embed_dim, bias=bias)
+        self.proj = nn.Linear(dim_feedforward, d_model, bias=bias)
         self.drop_prob = nn.Dropout(drop_prob)
         self.drop_prob1 = nn.Dropout(drop_prob)
 
@@ -476,7 +476,7 @@ class _FeedForward(nn.Module):
 
 
 class _TransformerEncoderLayer(nn.Module):
-    r"""Single Transformer encoder layer (pre-norm) combining MHSA and feed-forward.
+    """Single Transformer encoder layer (pre-norm) combining MHSA and feed-forward.
 
     The block follows the pattern:
     x <- x + MHSA(_LayerNorm(x))
@@ -485,27 +485,27 @@ class _TransformerEncoderLayer(nn.Module):
 
     def __init__(
         self,
-        embed_dim: int,
+        d_model: int,
         n_head: int,
         drop_prob: float = 0.0,
         dim_feedforward: Optional[int] = None,
         bias: bool = False,
-        activation: type[nn.Module] = nn.GELU,
+        activation: nn.Module = nn.GELU,
     ) -> None:
         super().__init__()
 
         if dim_feedforward is None:
-            dim_feedforward = 4 * embed_dim
+            dim_feedforward = 4 * d_model
             # note: preserve the original behaviour (print) from the provided code
             print(
-                "dim_feedforward is set to 4*embed_dim, the default in Vaswani et al. (Attention is all you need)"
+                "dim_feedforward is set to 4*d_model, the default in Vaswani et al. (Attention is all you need)"
             )
 
-        self.layer_norm_att = _LayerNorm(embed_dim, bias=bias)
-        self.mhsa = _MHSA(embed_dim, n_head, bias, drop_prob=drop_prob)
-        self.layer_norm_ff = _LayerNorm(embed_dim, bias=bias)
+        self.layer_norm_att = _LayerNorm(d_model, bias=bias)
+        self.mhsa = _MHSA(d_model, n_head, bias, drop_prob=drop_prob)
+        self.layer_norm_ff = _LayerNorm(d_model, bias=bias)
         self.feed_forward = _FeedForward(
-            embed_dim=embed_dim,
+            d_model=d_model,
             dim_feedforward=dim_feedforward,
             drop_prob=drop_prob,
             bias=bias,
@@ -518,7 +518,7 @@ class _TransformerEncoderLayer(nn.Module):
         Parameters
         ----------
         x : torch.Tensor
-            Input of shape (B, T, embed_dim).
+            Input of shape (B, T, d_model).
 
         Returns
         -------
@@ -531,13 +531,13 @@ class _TransformerEncoderLayer(nn.Module):
 
 
 class _TransformerEncoder(nn.Module):
-    r"""Stack of Transformer encoder layers.
+    """Stack of Transformer encoder layers.
 
     Parameters
     ----------
-    num_layers : int
+    n_blocks : int
         Number of encoder layers to stack.
-    embed_dim : int
+    d_model : int
         Dimensionality of embeddings.
     n_head : int
         Number of attention heads per layer.
@@ -549,26 +549,26 @@ class _TransformerEncoder(nn.Module):
 
     def __init__(
         self,
-        num_layers: int,
-        embed_dim: int,
+        n_blocks: int,
+        d_model: int,
         n_head: int,
         drop_prob: float,
         bias: bool,
-        activation: type[nn.Module] = nn.GELU,
+        activation: nn.Module = nn.GELU,
     ) -> None:
         super().__init__()
 
         self.encoder_block = nn.ModuleList(
             [
                 _TransformerEncoderLayer(
-                    embed_dim=embed_dim,
+                    d_model=d_model,
                     n_head=n_head,
                     drop_prob=drop_prob,
                     dim_feedforward=None,
                     bias=bias,
                     activation=activation,
                 )
-                for _ in range(num_layers)
+                for _ in range(n_blocks)
             ]
         )
 
@@ -576,7 +576,7 @@ class _TransformerEncoder(nn.Module):
         self.apply(self._init_weights)
         for pn, p in self.named_parameters():
             if pn.endswith("proj.weight"):
-                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * num_layers))
+                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_blocks))
 
     @staticmethod
     def _init_weights(module: nn.Module) -> None:
@@ -593,7 +593,7 @@ class _TransformerEncoder(nn.Module):
 
 
 class _Patcher(nn.Module):
-    r"""Patching encoding helper.
+    """Patching encoding helper.
 
     This module "patchifies" the original X entry in a ViT manner.
 

braindecode/models/sccnet.py

@@ -15,7 +15,7 @@ from braindecode.modules import LogActivation
 
 
 class SCCNet(EEGModuleMixin, nn.Module):
-    r"""SCCNet from Wei, C S (2019) [sccnet]_.
+    """SCCNet from Wei, C S (2019) [sccnet]_.
 
     :bdg-success:`Convolution`
 
@@ -155,7 +155,7 @@ class SCCNet(EEGModuleMixin, nn.Module):
         n_spatial_filters: int = 22,
         n_spatial_filters_smooth: int = 20,
         drop_prob: float = 0.5,
-        activation: type[nn.Module] = LogActivation,
+        activation: nn.Module = LogActivation,
         batch_norm_momentum: float = 0.1,
     ):
         super().__init__(

braindecode/models/shallow_fbcsp.py

@@ -2,8 +2,6 @@
 #
 # License: BSD (3-clause)
 
-from typing import Callable
-
 from einops.layers.torch import Rearrange
 from torch import nn
 from torch.nn import init
@@ -20,7 +18,7 @@ from braindecode.modules import (
 
 
 class ShallowFBCSPNet(EEGModuleMixin, nn.Sequential):
-    r"""Shallow ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
+    """Shallow ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
 
     :bdg-success:`Convolution`
 
@@ -83,9 +81,9 @@ class ShallowFBCSPNet(EEGModuleMixin, nn.Sequential):
         pool_time_length=75,
         pool_time_stride=15,
         final_conv_length="auto",
-        conv_nonlin: Callable = square,
+        conv_nonlin=square,
         pool_mode="mean",
-        activation_pool_nonlin: type[nn.Module] = SafeLog,
+        activation_pool_nonlin: nn.Module = SafeLog,
         split_first_layer=True,
         batch_norm=True,
         batch_norm_alpha=0.1,

braindecode/models/signal_jepa.py

@@ -24,7 +24,7 @@ _DEFAULT_CONV_LAYER_SPEC = (  # downsampling: 128Hz -> 1Hz, receptive field 1.18
 
 
 class _BaseSignalJEPA(EEGModuleMixin, nn.Module):
-    r"""Base class for the SignalJEPA models
+    """Base class for the SignalJEPA models
 
     Parameters
     ----------
@@ -144,9 +144,9 @@ class _BaseSignalJEPA(EEGModuleMixin, nn.Module):
 
 
 class SignalJEPA(_BaseSignalJEPA):
-    r"""Architecture introduced in signal-JEPA for self-supervised pre-training, Guetschel, P et al (2024) [1]_
+    """Architecture introduced in signal-JEPA for self-supervised pre-training, Guetschel, P et al (2024) [1]_
 
-    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Foundation Model`
+    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Large Brain Model`
 
     This model is not meant for classification but for SSL pre-training.
     Its output shape depends on the input shape.
@@ -232,9 +232,9 @@ class SignalJEPA(_BaseSignalJEPA):
 
 
 class SignalJEPA_Contextual(_BaseSignalJEPA):
-    r"""Contextual downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
+    """Contextual downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
 
-    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Foundation Model`
+    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Large Brain Model`
 
     This architecture is one of the variants of :class:`SignalJEPA`
     that can be used for classification purposes.
@@ -405,9 +405,9 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
 
 
 class SignalJEPA_PostLocal(_BaseSignalJEPA):
-    r"""Post-local downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
+    """Post-local downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
 
-    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Foundation Model`
+    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Large Brain Model`
 
     This architecture is one of the variants of :class:`SignalJEPA`
     that can be used for classification purposes.
@@ -556,9 +556,9 @@ class SignalJEPA_PostLocal(_BaseSignalJEPA):
 
 
 class SignalJEPA_PreLocal(_BaseSignalJEPA):
-    r"""Pre-local downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
+    """Pre-local downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
 
-    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Foundation Model`
+    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Large Brain Model`
 
     This architecture is one of the variants of :class:`SignalJEPA`
     that can be used for classification purposes.
@@ -569,33 +569,6 @@ class SignalJEPA_PreLocal(_BaseSignalJEPA):
 
     .. versionadded:: 0.9
 
-    .. 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 SignalJEPA_PreLocal
-
-           # Load pre-trained model from Hugging Face Hub
-           model = SignalJEPA_PreLocal.from_pretrained(
-               "braindecode/SignalJEPA-PreLocal-pretrained"
-           )
-
-       To push your own trained model to the Hub:
-
-       .. code-block:: python
-
-           # After training your model
-           model.push_to_hub(
-               repo_id="username/my-sjepa-model",
-               commit_message="Upload trained SignalJEPA model",
-           )
-
-       Requires installing ``braindecode[hug]`` for Hub integration.
-
     Parameters
     ----------
     n_spat_filters : int
@@ -745,7 +718,7 @@ class SignalJEPA_PreLocal(_BaseSignalJEPA):
 
 
 class _ConvFeatureEncoder(nn.Sequential):
-    r"""Convolutional feature encoder for EEG data.
+    """Convolutional feature encoder for EEG data.
 
     Computes successive 1D convolutions (with activations) over the time
     dimension of the input EEG signal.
@@ -865,7 +838,7 @@ class _ConvFeatureEncoder(nn.Sequential):
 
 
 class _ChannelEmbedding(nn.Embedding):
-    r"""Embedding layer for EEG channels.
+    """Embedding layer for EEG channels.
 
     The difference with a regular :class:`nn.Embedding` is that the embedding
     vectors are initialized with a positional encodding of the channel locations.
@@ -926,7 +899,7 @@ class _ChannelEmbedding(nn.Embedding):
 
 
 class _PosEncoder(nn.Module):
-    r"""Positional encoder for EEG data.
+    """Positional encoder for EEG data.
 
     Parameters
     ----------

braindecode/models/sinc_shallow.py

@@ -1,4 +1,5 @@
 import math
+from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -9,7 +10,7 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class SincShallowNet(EEGModuleMixin, nn.Module):
-    r"""Sinc-ShallowNet from Borra, D et al (2020) [borra2020]_.
+    """Sinc-ShallowNet from Borra, D et al (2020) [borra2020]_.
 
     :bdg-success:`Convolution` :bdg-warning:`Interpretability`
 
@@ -92,7 +93,7 @@ class SincShallowNet(EEGModuleMixin, nn.Module):
         num_time_filters: int = 32,
         time_filter_len: int = 33,
         depth_multiplier: int = 2,
-        activation: type[nn.Module] | None = nn.ELU,
+        activation: Optional[nn.Module] = nn.ELU,
         drop_prob: float = 0.5,
         first_freq: float = 5.0,
         min_freq: float = 1.0,
@@ -210,7 +211,7 @@ class SincShallowNet(EEGModuleMixin, nn.Module):
 
 
 class _SincFilter(nn.Module):
-    r"""Sinc-Based Convolutional Layer for Band-Pass Filtering from Ravanelli and Bengio (2018) [ravanelli]_.
+    """Sinc-Based Convolutional Layer for Band-Pass Filtering from Ravanelli and Bengio (2018) [ravanelli]_.
 
     The `SincFilter` layer implements a convolutional layer where each kernel is
     defined using a parametrized sinc function.

braindecode/models/sleep_stager_blanco_2020.py

@@ -9,7 +9,7 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
-    r"""Sleep staging architecture from Blanco et al (2020) from [Blanco2020]_
+    """Sleep staging architecture from Blanco et al. (2020) from [Blanco2020]_
 
     :bdg-success:`Convolution`
 
@@ -68,7 +68,7 @@ class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
         drop_prob=0.5,
         apply_batch_norm=False,
         return_feats=False,
-        activation: type[nn.Module] = nn.ReLU,
+        activation: nn.Module = nn.ReLU,
         chs_info=None,
         n_times=None,
     ):

braindecode/models/sleep_stager_chambon_2018.py

@@ -11,7 +11,7 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
-    r"""Sleep staging architecture from Chambon et al. (2018) [Chambon2018]_.
+    """Sleep staging architecture from Chambon et al. (2018) [Chambon2018]_.
 
     :bdg-success:`Convolution`
 
@@ -70,7 +70,7 @@ class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
         time_conv_size_s=0.5,
         max_pool_size_s=0.125,
         pad_size_s=0.25,
-        activation: type[nn.Module] = nn.ReLU,
+        activation: nn.Module = nn.ReLU,
         input_window_seconds=None,
         n_outputs=5,
         drop_prob=0.25,

braindecode/models/sparcnet.py

@@ -11,7 +11,7 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class SPARCNet(EEGModuleMixin, nn.Module):
-    r"""Seizures, Periodic and Rhythmic pattern Continuum Neural Network (SPaRCNet) from Jing et al (2023) [jing2023]_.
+    """Seizures, Periodic and Rhythmic pattern Continuum Neural Network (SPaRCNet) from Jing et al. (2023) [jing2023]_.
 
     :bdg-success:`Convolution`
 
@@ -69,7 +69,7 @@ class SPARCNet(EEGModuleMixin, nn.Module):
         drop_prob: float = 0.5,
         conv_bias: bool = True,
         batch_norm: bool = True,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
         kernel_size_conv0: int = 7,
         kernel_size_conv1: int = 1,
         kernel_size_conv2: int = 3,
@@ -213,7 +213,7 @@ class SPARCNet(EEGModuleMixin, nn.Module):
 
 
 class _DenseLayer(nn.Sequential):
-    r"""
+    """
     A densely connected layer with batch normalization and dropout.
 
     Parameters
@@ -252,7 +252,7 @@ class _DenseLayer(nn.Sequential):
         drop_prob: float = 0.5,
         conv_bias: bool = True,
         batch_norm: bool = True,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
         kernel_size_conv1: int = 1,
         kernel_size_conv2: int = 3,
         stride_conv1: int = 1,
@@ -302,7 +302,7 @@ class _DenseLayer(nn.Sequential):
 
 
 class _DenseBlock(nn.Sequential):
-    r"""
+    """
     A densely connected block that uses DenseLayers.
 
     Parameters
@@ -344,7 +344,7 @@ class _DenseBlock(nn.Sequential):
         drop_prob=0.5,
         conv_bias=True,
         batch_norm=True,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
         kernel_size_conv1: int = 1,
         kernel_size_conv2: int = 3,
         stride_conv1: int = 1,
@@ -371,7 +371,7 @@ class _DenseBlock(nn.Sequential):
 
 
 class _TransitionLayer(nn.Sequential):
-    r"""
+    """
     A pooling transition layer.
 
     Parameters
@@ -403,7 +403,7 @@ class _TransitionLayer(nn.Sequential):
         out_channels,
         conv_bias=True,
         batch_norm=True,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
         kernel_size_trans: int = 2,
         stride_trans: int = 2,
     ):