braindecode 1.3.0.dev175435903__tar.gz → 1.3.0.dev177628147__tar.gz

{braindecode-1.3.0.dev175435903/braindecode.egg-info → braindecode-1.3.0.dev177628147}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.0.dev175435903
+Version: 1.3.0.dev177628147
 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>
@@ -52,9 +52,6 @@ Requires-Dist: pytest-cov; extra == "tests"
 Requires-Dist: codecov; extra == "tests"
 Requires-Dist: pytest_cases; extra == "tests"
 Requires-Dist: mypy; extra == "tests"
-Provides-Extra: typing
-Requires-Dist: pydantic<3.0,>=2.0; extra == "typing"
-Requires-Dist: numpydantic>=1.7; extra == "typing"
 Provides-Extra: docs
 Requires-Dist: sphinx_gallery; extra == "docs"
 Requires-Dist: sphinx_rtd_theme; extra == "docs"
@@ -79,7 +76,6 @@ Requires-Dist: braindecode[hub]; extra == "all"
 Requires-Dist: braindecode[tests]; extra == "all"
 Requires-Dist: braindecode[docs]; extra == "all"
 Requires-Dist: braindecode[eegprep]; extra == "all"
-Requires-Dist: braindecode[typing]; extra == "all"
 Dynamic: license-file
 
 .. image:: https://badges.gitter.im/braindecodechat/community.svg

{braindecode-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/braindecode/datautil/serialization.py

@@ -138,7 +138,7 @@ def _load_signals(fif_file, preload, is_raw):
         with open(pkl_file, "rb") as f:
             signals = pickle.load(f)
 
-        if all(Path(f).exists() for f in signals.filenames):
+        if all(f.exists() for f in signals.filenames):
             if preload:
                 signals.load_data()
             return signals

{braindecode-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/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,13 +248,13 @@ 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,
         tcn_drop_prob=0.3,
-        tcn_activation: type[nn.Module] = nn.ELU,
+        tcn_activation: nn.Module = nn.ELU,
         concat=False,
         max_norm_const=0.25,
         chs_info=None,
@@ -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)
@@ -656,7 +656,7 @@ class _TCNResidualBlock(nn.Module):
         kernel_size=4,
         n_filters=32,
         dropout=0.3,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
         dilation=1,
     ):
         super().__init__()

{braindecode-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/braindecode/models/attentionbasenet.py

@@ -235,7 +235,7 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
     kernel_size : int, default=9
         The kernel size used in certain types of attention mechanisms for convolution
         operations.
-    activation: type[nn.Module] = nn.ELU,
+    activation: nn.Module, default=nn.ELU
         Activation function class to apply. Should be a PyTorch activation
         module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
     extra_params : bool, default=False
@@ -277,7 +277,7 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         freq_idx: int = 0,
         n_codewords: int = 4,
         kernel_size: int = 9,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
         extra_params: bool = False,
     ):
         super(AttentionBaseNet, self).__init__()
@@ -453,7 +453,7 @@ class _FeatureExtractor(nn.Module):
         pool_length: int = 75,
         pool_stride: int = 15,
         drop_prob: float = 0.5,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
     ):
         super().__init__()
 
@@ -592,7 +592,7 @@ class _ChannelAttentionBlock(nn.Module):
         n_codewords: int = 4,
         kernel_size: int = 9,
         extra_params: bool = False,
-        activation: type[nn.Module] = nn.ELU,
+        activation: nn.Module = nn.ELU,
     ):
         super().__init__()
         self.conv = nn.Sequential(

{braindecode-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/braindecode/models/attn_sleep.py

@@ -90,8 +90,8 @@ class AttnSleep(EEGModuleMixin, nn.Module):
         d_ff=120,
         n_attn_heads=5,
         drop_prob=0.1,
-        activation_mrcnn: type[nn.Module] = nn.GELU,
-        activation: type[nn.Module] = nn.ReLU,
+        activation_mrcnn: nn.Module = nn.GELU,
+        activation: nn.Module = nn.ReLU,
         input_window_seconds=None,
         n_outputs=None,
         after_reduced_cnn_size=30,
@@ -230,7 +230,7 @@ class _SEBasicBlock(nn.Module):
         planes,
         stride=1,
         downsample=None,
-        activation: type[nn.Module] = nn.ReLU,
+        activation: nn.Module = nn.ReLU,
         *,
         reduction=16,
     ):
@@ -278,8 +278,8 @@ class _MRCNN(nn.Module):
         self,
         after_reduced_cnn_size,
         kernel_size=7,
-        activation: type[nn.Module] = nn.GELU,
-        activation_se: type[nn.Module] = nn.ReLU,
+        activation: nn.Module = nn.GELU,
+        activation_se: nn.Module = nn.ReLU,
     ):
         super(_MRCNN, self).__init__()
         drate = 0.5
@@ -325,7 +325,7 @@ class _MRCNN(nn.Module):
         )
 
     def _make_layer(
-        self, block, planes, blocks, stride=1, activate: type[nn.Module] = nn.ReLU
+        self, block, planes, blocks, stride=1, activate: nn.Module = nn.ReLU
     ):  # makes residual SE block
         downsample = None
         if stride != 1 or self.inplanes != planes * block.expansion:
@@ -526,9 +526,7 @@ class _EncoderLayer(nn.Module):
 class _PositionwiseFeedForward(nn.Module):
     """Positionwise feed-forward network."""
 
-    def __init__(
-        self, d_model, d_ff, dropout=0.1, activation: type[nn.Module] = nn.ReLU
-    ):
+    def __init__(self, d_model, d_ff, dropout=0.1, activation: nn.Module = nn.ReLU):
         super().__init__()
         self.w_1 = nn.Linear(d_model, d_ff)
         self.w_2 = nn.Linear(d_ff, d_model)

{braindecode-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/braindecode/models/bendr.py

@@ -176,7 +176,7 @@ class BENDR(EEGModuleMixin, nn.Module):
         projection_head=False,  # Whether encoder should project back to input feature size (unused in original fine-tuning)
         drop_prob=0.1,  # General dropout probability (paper: 0.15 for pretraining, 0.0 for fine-tuning)
         layer_drop=0.0,  # Probability of dropping transformer layers during training (paper: 0.01 for pretraining)
-        activation: type[nn.Module] = nn.GELU,  # Activation function
+        activation=nn.GELU,  # Activation function
         # Transformer specific parameters
         transformer_layers=8,
         transformer_heads=8,

{braindecode-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/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,
@@ -87,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,
@@ -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,
         )
@@ -250,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)
@@ -347,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
@@ -360,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
@@ -380,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.dev175435903 → braindecode-1.3.0.dev177628147}/braindecode/models/contrawr.py

@@ -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,
@@ -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(

{braindecode-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/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
@@ -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-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/braindecode/models/deep4.py

@@ -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-1.3.0.dev175435903 → braindecode-1.3.0.dev177628147}/braindecode/models/deepsleepnet.py

@@ -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__(
@@ -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(
@@ -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(