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

braindecode/datasets/mne.py

@@ -9,7 +9,7 @@ import mne
 import numpy as np
 import pandas as pd
 
-from .base import BaseConcatDataset, BaseDataset, WindowsDataset
+from .base import BaseConcatDataset, RawDataset, WindowsDataset
 
 
 def create_from_mne_raw(
@@ -75,11 +75,9 @@ def create_from_mne_raw(
                 f"length of 'raws' ({len(raws)}) and 'description' "
                 f"({len(descriptions)}) has to match"
             )
-        base_datasets = [
-            BaseDataset(raw, desc) for raw, desc in zip(raws, descriptions)
-        ]
+        base_datasets = [RawDataset(raw, desc) for raw, desc in zip(raws, descriptions)]
     else:
-        base_datasets = [BaseDataset(raw) for raw in raws]
+        base_datasets = [RawDataset(raw) for raw in raws]
 
     base_datasets = BaseConcatDataset(base_datasets)
     windows_datasets = create_windows_from_events(

braindecode/datasets/moabb.py

@@ -18,7 +18,7 @@ import pandas as pd
 
 from braindecode.util import _update_moabb_docstring
 
-from .base import BaseConcatDataset, BaseDataset
+from .base import BaseConcatDataset, RawDataset
 
 
 def _find_dataset_in_moabb(dataset_name, dataset_kwargs=None):
@@ -164,7 +164,7 @@ class MOABBDataset(BaseConcatDataset):
             dataset_load_kwargs=dataset_load_kwargs,
         )
         all_base_ds = [
-            BaseDataset(raw, row) for raw, (_, row) in zip(raws, description.iterrows())
+            RawDataset(raw, row) for raw, (_, row) in zip(raws, description.iterrows())
         ]
         super().__init__(all_base_ds)
 

braindecode/datasets/nmt.py

@@ -31,7 +31,7 @@ import pandas as pd
 from joblib import Parallel, delayed
 from mne.datasets import fetch_dataset
 
-from braindecode.datasets.base import BaseConcatDataset, BaseDataset
+from braindecode.datasets.base import BaseConcatDataset, RawDataset
 
 NMT_URL = "https://zenodo.org/record/10909103/files/NMT.zip"
 NMT_archive_name = "NMT.zip"
@@ -172,7 +172,7 @@ class NMT(BaseConcatDataset):
         d["n_samples"] = raw.n_times
         d["sfreq"] = raw.info["sfreq"]
         d["train"] = "train" in d.path.split(os.sep)
-        base_dataset = BaseDataset(raw, d, target_name)
+        base_dataset = RawDataset(raw, d, target_name)
         return base_dataset
 
 

braindecode/datasets/sleep_physio_challe_18.py

@@ -21,8 +21,7 @@ from mne.datasets.sleep_physionet._utils import _fetch_one
 from mne.datasets.utils import _get_path
 from mne.utils import warn
 
-from braindecode.datasets import BaseConcatDataset, BaseDataset
-from braindecode.preprocessing.preprocess import _preprocess
+from braindecode.datasets import BaseConcatDataset, RawDataset
 
 PC18_DIR = op.join(op.dirname(__file__), "data", "pc18")
 PC18_RECORDS = op.join(PC18_DIR, "sleep_records.csv")
@@ -404,9 +403,11 @@ class SleepPhysionetChallenge2018(BaseConcatDataset):
             },
             name="",
         )
-        base_dataset = BaseDataset(raw_file, desc)
+        base_dataset = RawDataset(raw_file, desc)
 
         if preproc is not None:
+            from braindecode.preprocessing.preprocess import _preprocess
+
             _preprocess(base_dataset, None, preproc)
 
         return base_dataset

braindecode/datasets/sleep_physionet.py

@@ -12,7 +12,7 @@ import numpy as np
 import pandas as pd
 from mne.datasets.sleep_physionet.age import fetch_data
 
-from .base import BaseConcatDataset, BaseDataset
+from .base import BaseConcatDataset, RawDataset
 
 
 class SleepPhysionet(BaseConcatDataset):
@@ -71,7 +71,7 @@ class SleepPhysionet(BaseConcatDataset):
                 crop_wake_mins=crop_wake_mins,
                 crop=crop,
             )
-            base_ds = BaseDataset(raw, desc)
+            base_ds = RawDataset(raw, desc)
             all_base_ds.append(base_ds)
         super().__init__(all_base_ds)
 

braindecode/datasets/tuh.py

@@ -22,7 +22,7 @@ import numpy as np
 import pandas as pd
 from joblib import Parallel, delayed
 
-from .base import BaseConcatDataset, BaseDataset
+from .base import BaseConcatDataset, RawDataset
 
 
 class TUH(BaseConcatDataset):
@@ -214,7 +214,7 @@ class TUH(BaseConcatDataset):
             d["report"] = physician_report
         additional_description = pd.Series(d)
         description = pd.concat([description, additional_description])
-        base_dataset = BaseDataset(raw, description, target_name=target_name)
+        base_dataset = RawDataset(raw, description, target_name=target_name)
         return base_dataset
 
 

braindecode/datasets/xy.py

@@ -12,7 +12,7 @@ import numpy as np
 import pandas as pd
 from numpy.typing import ArrayLike, NDArray
 
-from .base import BaseConcatDataset, BaseDataset
+from .base import BaseConcatDataset, RawDataset
 
 log = logging.getLogger(__name__)
 
@@ -69,7 +69,7 @@ def create_from_X_y(
         n_samples_per_x.append(x.shape[1])
         info = mne.create_info(ch_names=ch_names, sfreq=sfreq)
         raw = mne.io.RawArray(x, info)
-        base_dataset = BaseDataset(
+        base_dataset = RawDataset(
             raw, pd.Series({"target": target}), target_name="target"
         )
         base_datasets.append(base_dataset)

braindecode/datautil/serialization.py

@@ -19,8 +19,8 @@ from joblib import Parallel, delayed
 
 from ..datasets.base import (
     BaseConcatDataset,
-    BaseDataset,
     EEGWindowsDataset,
+    RawDataset,
     WindowsDataset,
 )
 
@@ -35,7 +35,7 @@ def save_concat_dataset(path, concat_dataset, overwrite=False):
 
 
 def _outdated_load_concat_dataset(path, preload, ids_to_load=None, target_name=None):
-    """Load a stored BaseConcatDataset of BaseDatasets or WindowsDatasets from
+    """Load a stored BaseConcatDataset from
     files.
 
     Parameters
@@ -52,7 +52,7 @@ def _outdated_load_concat_dataset(path, preload, ids_to_load=None, target_name=N
 
     Returns
     -------
-    concat_dataset: BaseConcatDataset of BaseDatasets or WindowsDatasets
+    concat_dataset: BaseConcatDataset
     """
     # assume we have a single concat dataset to load
     is_raw = (path / "0-raw.fif").is_file()
@@ -87,7 +87,7 @@ def _outdated_load_concat_dataset(path, preload, ids_to_load=None, target_name=N
         for i_signal, signal in enumerate(all_signals):
             if is_raw:
                 datasets.append(
-                    BaseDataset(
+                    RawDataset(
                         signal, description.iloc[i_signal], target_name=target_name
                     )
                 )
@@ -138,12 +138,17 @@ def _load_signals(fif_file, preload, is_raw):
         with open(pkl_file, "rb") as f:
             signals = pickle.load(f)
 
-        # If the file has been moved together with the pickle file, make sure
-        # the path links to correct fif file.
-        signals._fname = str(fif_file)
-        if preload:
-            signals.load_data()
-        return signals
+        if all(f.exists() for f in signals.filenames):
+            if preload:
+                signals.load_data()
+            return signals
+        else:  # This may happen if the file has been moved together with the pickle file.
+            warnings.warn(
+                f"Pickle file {pkl_file} exists, but the referenced fif "
+                "file(s) do not exist. Will read the fif file(s) directly "
+                "and re-create the pickle file.",
+                UserWarning,
+            )
 
     # If pickle didn't exist read via mne (likely slower) and save pkl after
     if is_raw:
@@ -170,7 +175,7 @@ def _load_signals(fif_file, preload, is_raw):
 
 
 def load_concat_dataset(path, preload, ids_to_load=None, target_name=None, n_jobs=1):
-    """Load a stored BaseConcatDataset of BaseDatasets or WindowsDatasets from
+    """Load a stored BaseConcatDataset from
     files.
 
     Parameters
@@ -189,7 +194,7 @@ def load_concat_dataset(path, preload, ids_to_load=None, target_name=None, n_job
 
     Returns
     -------
-    concat_dataset: BaseConcatDataset of BaseDatasets or WindowsDatasets
+    concat_dataset: BaseConcatDataset
     """
     # Make sure we always work with a pathlib.Path
     path = Path(path)
@@ -261,7 +266,7 @@ def _load_parallel(path, i, preload, is_raw, has_stored_windows):
         target_name = json.load(open(target_file_path, "r"))["target_name"]
 
     if is_raw and (not has_stored_windows):
-        dataset = BaseDataset(signals, description, target_name)
+        dataset = RawDataset(signals, description, target_name)
     else:
         window_kwargs = _load_kwargs_json("window_kwargs", sub_dir)
         windows_ds_kwargs = [

braindecode/eegneuralnet.py

@@ -189,6 +189,8 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
                 "Skipping setting signal-related parameters from data."
             )
             return
+        if classes is None:
+            classes = getattr(self, "classes", None)
         # get kwargs from signal:
         signal_kwargs = dict()
         # Using shape to work both with torch.tensor and numpy.array:

braindecode/functional/functions.py

@@ -181,20 +181,24 @@ def plv_time(x, forward_fourier=True, epsilon: float = 1e-6):
     The Phase Locking Value (PLV) is a measure of the synchronization between
     different channels by evaluating the consistency of phase differences
     over time. It ranges from 0 (no synchronization) to 1 (perfect
-    synchronization) [1]_.
+    synchronization) [Lachaux1999]_.
 
     Parameters
     ----------
     x : torch.Tensor
         Input tensor containing the signal data.
+
         - If `forward_fourier` is `True`, the shape should be `(..., channels, time)`.
         - If `forward_fourier` is `False`, the shape should be `(..., channels, freqs, 2)`,
           where the last dimension represents the real and imaginary parts.
+
     forward_fourier : bool, optional
         Specifies the format of the input tensor `x`.
+
         - If `True`, `x` is assumed to be in the time domain.
         - If `False`, `x` is assumed to be in the Fourier domain with separate real and
           imaginary components.
+
         Default is `True`.
     epsilon : float, default 1e-6
         Small numerical value to ensure positivity constraint on the complex part
@@ -207,7 +211,7 @@ def plv_time(x, forward_fourier=True, epsilon: float = 1e-6):
 
     References
     ----------
-    [1] Lachaux, J. P., Rodriguez, E., Martinerie, J., & Varela, F. J. (1999).
+    .. [Lachaux1999] Lachaux, J. P., Rodriguez, E., Martinerie, J., & Varela, F. J. (1999).
         Measuring phase synchrony in brain signals. Human brain mapping,
         8(4), 194-208.
     """

braindecode/functional/initialization.py

@@ -5,9 +5,8 @@ from torch import nn
 
 def glorot_weight_zero_bias(model):
     """Initialize parameters of all modules by initializing weights with
-    glorot
-     uniform/xavier initialization, and setting biases to zero. Weights from
-     batch norm layers are set to 1.
+    glorot uniform/xavier initialization, and setting biases to zero. Weights from
+    batch norm layers are set to 1.
 
     Parameters
     ----------

braindecode/models/__init__.py

@@ -4,7 +4,9 @@ Some predefined network architectures for EEG decoding.
 
 from .atcnet import ATCNet
 from .attentionbasenet import AttentionBaseNet
+from .attn_sleep import AttnSleep
 from .base import EEGModuleMixin
+from .bendr import BENDR
 from .biot import BIOT
 from .contrawr import ContraWR
 from .ctnet import CTNet
@@ -15,9 +17,8 @@ from .eeginception_erp import EEGInceptionERP
 from .eeginception_mi import EEGInceptionMI
 from .eegitnet import EEGITNet
 from .eegminer import EEGMiner
-from .eegnet import EEGNetv1, EEGNetv4
+from .eegnet import EEGNet, EEGNetv4
 from .eegnex import EEGNeX
-from .eegresnet import EEGResNet
 from .eegsimpleconv import EEGSimpleConv
 from .eegtcnet import EEGTCNet
 from .fbcnet import FBCNet
@@ -27,6 +28,7 @@ from .hybrid import HybridNet
 from .ifnet import IFNet
 from .labram import Labram
 from .msvtnet import MSVTNet
+from .patchedtransformer import PBT
 from .sccnet import SCCNet
 from .shallow_fbcsp import ShallowFBCSPNet
 from .signal_jepa import (
@@ -38,12 +40,12 @@ from .signal_jepa import (
 from .sinc_shallow import SincShallowNet
 from .sleep_stager_blanco_2020 import SleepStagerBlanco2020
 from .sleep_stager_chambon_2018 import SleepStagerChambon2018
-from .sleep_stager_eldele_2021 import SleepStagerEldele2021
 from .sparcnet import SPARCNet
+from .sstdpn import SSTDPN
 from .syncnet import SyncNet
 from .tcn import BDTCN, TCN
 from .tidnet import TIDNet
-from .tsinception import TSceptionV1
+from .tsinception import TSception
 from .usleep import USleep
 from .util import _init_models_dict, models_mandatory_parameters
 
@@ -53,9 +55,11 @@ _init_models_dict()
 
 __all__ = [
     "ATCNet",
+    "AttnSleep",
     "AttentionBaseNet",
     "EEGModuleMixin",
     "BIOT",
+    "BENDR",
     "ContraWR",
     "CTNet",
     "Deep4Net",
@@ -65,10 +69,9 @@ __all__ = [
     "EEGInceptionMI",
     "EEGITNet",
     "EEGMiner",
-    "EEGNetv1",
+    "EEGNet",
     "EEGNetv4",
     "EEGNeX",
-    "EEGResNet",
     "EEGSimpleConv",
     "EEGTCNet",
     "FBCNet",
@@ -78,6 +81,7 @@ __all__ = [
     "IFNet",
     "Labram",
     "MSVTNet",
+    "PBT",
     "SCCNet",
     "ShallowFBCSPNet",
     "SignalJEPA",
@@ -85,15 +89,15 @@ __all__ = [
     "SignalJEPA_PostLocal",
     "SignalJEPA_PreLocal",
     "SincShallowNet",
+    "SSTDPN",
     "SleepStagerBlanco2020",
     "SleepStagerChambon2018",
-    "SleepStagerEldele2021",
     "SPARCNet",
     "SyncNet",
     "BDTCN",
     "TCN",
     "TIDNet",
-    "TSceptionV1",
+    "TSception",
     "USleep",
     "_init_models_dict",
     "models_mandatory_parameters",

braindecode/models/atcnet.py

@@ -13,13 +13,153 @@ from braindecode.modules import CausalConv1d, Ensure4d, MaxNormLinear
 
 
 class ATCNet(EEGModuleMixin, nn.Module):
-    """ATCNet model from Altaheri et al. (2022) [1]_
+    """ATCNet from Altaheri et al. (2022) [1]_.
 
-    Pytorch implementation based on official tensorflow code [2]_.
+    :bdg-success:`Convolution` :bdg-secondary:`Recurrent` :bdg-info:`Small Attention`
 
     .. figure:: https://user-images.githubusercontent.com/25565236/185449791-e8539453-d4fa-41e1-865a-2cf7e91f60ef.png
-       :align: center
-       :alt: ATCNet Architecture
+        :align: center
+        :alt: ATCNet Architecture
+        :width: 650px
+
+    .. rubric:: Architectural Overview
+
+    ATCNet is a *convolution-first* architecture augmented with a *lightweight attention–TCN*
+    sequence module. The end-to-end flow is:
+
+    - (i) :class:`_ConvBlock` learns temporal filter-banks and spatial projections (EEGNet-style),
+      downsampling time to a compact feature map;
+
+    - (ii) Sliding Windows carve overlapping temporal windows from this map;
+
+    - (iii) for each window, :class:`_AttentionBlock` applies small multi-head self-attention
+      over time, followed by a :class:`_TCNResidualBlock` stack (causal, dilated);
+
+    - (iv) window-level features are aggregated (mean of window logits or concatenation)
+      and mapped via a max-norm–constrained linear layer.
+
+    Relative to ViT, ATCNet replaces linear patch projection with learned *temporal–spatial*
+    convolutions; it processes *parallel* window encoders (attention→TCN) instead of a deep
+    stack; and swaps the MLP head for a TCN suited to 1-D EEG sequences.
+
+    .. rubric:: Macro Components
+
+    - :class:`_ConvBlock` **(Shallow conv stem → feature map)**
+
+        - *Operations.*
+        - **Temporal conv** (:class:`torch.nn.Conv2d`) with kernel ``(L_t, 1)`` builds a
+            FIR-like filter bank (``F1`` maps).
+        - **Depthwise spatial conv** (:class:`torch.nn.Conv2d`, ``groups=F1``) with kernel
+          ``(1, n_chans)`` learns per-filter spatial projections (akin to EEGNet's CSP-like step).
+        - **BN → ELU → AvgPool → Dropout** to stabilize and condense activations.
+        - **Refining temporal conv** (:class:`torch.nn.Conv2d`) with kernel ``(L_r, 1)`` +
+          **BN → ELU → AvgPool → Dropout**.
+
+    The output shape is ``(B, F2, T_c, 1)`` with ``F2 = F1·D`` and ``T_c = T/(P1·P2)``.
+    Temporal kernels behave as FIR filters; the depthwise-spatial conv yields frequency-specific
+    topographies. Pooling acts as a local integrator, reducing variance and imposing a
+    useful inductive bias on short EEG windows.
+
+    - **Sliding-Window Sequencer**
+
+        From the condensed time axis (length ``T_c``), ATCNet forms ``n`` overlapping windows
+        of width ``T_w = T_c - n + 1`` (one start per index). Each window produces a sequence
+        ``(B, F2, T_w)`` forwarded to its own attention-TCN branch. This creates *parallel*
+        encoders over shifted contexts and is key to robustness on nonstationary EEG.
+
+    - :class:`_AttentionBlock` **(small MHA on temporal positions)**
+
+        Attention here is *local to a window* and purely temporal.
+
+        - *Operations.*
+        - Rearrange to ``(B, T_w, F2)``,
+        - Normalization :class:`torch.nn.LayerNorm`
+        - Custom MultiHeadAttention :class:`_MHA` (``num_heads=H``, per-head dim ``d_h``) + residual add,
+        - Dropout :class:`torch.nn.Dropout`
+        - Rearrange back to ``(B, F2, T_w)``.
+
+        *Role.* Re-weights evidence across the window, letting the model emphasize informative
+        segments (onsets, bursts) before causal convolutions aggregate history.
+
+    - :class:`_TCNResidualBlock` **(causal dilated temporal CNN)**
+
+        - *Operations.*
+        - Two :class:`braindecode.modules.CausalConv1d` layers per block with dilation  ``1, 2, 4, …``
+        - Across blocks of `torch.nn.ELU` + `torch.nn.BatchNorm1d` + `torch.nn.Dropout`) +
+          a residual (identity or 1x1 mapping).
+        - The final feature used per window is the *last* causal step ``[..., -1]`` (forecast-style).
+
+        *Role.* Efficient long-range temporal integration with stable gradients; the dilated
+        receptive field complements attention's soft selection.
+
+    - **Aggregation & Classifier**
+
+        - *Operations.*
+        - Either (a) map each window feature ``(B, F2)`` to logits via :class:`braindecode.modules.MaxNormLinear`
+        and **average** across windows (default, matching official code), or
+        - (b) **concatenate** all window features ``(B, n·F2)`` and apply a single :class:`MaxNormLinear`.
+        The max-norm constraint regularizes the readout.
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal.** Temporal structure is learned in three places:
+        - (1) the stem's wide ``(L_t, 1)`` conv (learned filter bank),
+        - (2) the refining ``(L_r, 1)`` conv after pooling (short-term dynamics), and
+        - (3) the TCN's causal 1-D convolutions with exponentially increasing dilation
+          (long-range dependencies). The minimum sequence length required by the TCN stack is
+          ``(K_t - 1)·2^{L-1} + 1``; the implementation *auto-scales* kernels/pools/windows
+          when inputs are shorter to preserve feasibility.
+
+    - **Spatial.** A depthwise spatial conv spans the **full montage** (kernel ``(1, n_chans)``),
+        producing *per-temporal-filter* spatial projections (no cross-filter mixing at this step).
+        This mirrors EEGNet's interpretability: each temporal filter has its own spatial pattern.
+
+
+    .. rubric:: Attention / Sequential Modules
+
+    - **Type.** Multi-head self-attention with ``H`` heads and per-head dim ``d_h`` implemented
+      in :class:`_MHA`, allowing ``embed_dim = H·d_h`` independent of input and output dims.
+    - **Shapes.** ``(B, F2, T_w) → (B, T_w, F2) → (B, F2, T_w)``. Attention operates along
+      the **temporal** axis within a window; channels/features stay in the embedding dim ``F2``.
+    - **Role.** Highlights salient temporal positions prior to causal convolution; small attention
+      keeps compute modest while improving context modeling over pooled features.
+
+    .. rubric:: Additional Mechanisms
+
+    - **Parallel encoders over shifted windows.** Improves montage/phase robustness by
+      ensembling nearby contexts rather than committing to a single segmentation.
+    - **Max-norm classifier.** Enforces weight norm constraints at the readout, a common
+      stabilization trick in EEG decoding.
+    - **ViT vs. ATCNet (design choices).** Convolutional *nonlinear* projection rather than
+      linear patchification; attention followed by **TCN** (not MLP); *parallel* window
+      encoders rather than stacked encoders.
+
+    .. rubric:: Usage and Configuration
+
+    - ``conv_block_n_filters (F1)``, ``conv_block_depth_mult (D)`` → capacity of the stem
+      (with ``F2 = F1·D`` feeding attention/TCN), dimensions aligned to ``F2``, like :class:`EEGNet`.
+    - 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).
+    - ``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.
+    - **Aggregation choice.** ``concat=False`` (default, average of per-window logits) matches
+      the official code; ``concat=True`` mirrors the paper's concatenation variant.
+
+
+    Notes
+    -----
+    - Inputs substantially shorter than the implied minimum length trigger **automatic
+      downscaling** of kernels, pools, windows, and TCN kernel size to maintain validity.
+    - The attention–TCN sequence operates **per window**; the last causal step is used as the
+      window feature, aligning the temporal semantics across windows.
+
+    .. versionadded:: 1.1
+
+        - More detailed documentation of the model.
+
 
     Parameters
     ----------
@@ -85,15 +225,13 @@ class ATCNet(EEGModuleMixin, nn.Module):
         Maximum L2-norm constraint imposed on weights of the last
         fully-connected layer. Defaults to 0.25.
 
-
     References
     ----------
-    .. [1] H. Altaheri, G. Muhammad and M. Alsulaiman,
-        Physics-informed attention temporal convolutional network for EEG-based
-        motor imagery classification in IEEE Transactions on Industrial Informatics,
-        2022, doi: 10.1109/TII.2022.3197419.
-    .. [2] EEE-ATCNet implementation.
-       https://github.com/Altaheri/EEG-ATCNet/blob/main/models.py
+    .. [1] H. Altaheri, G. Muhammad, M. Alsulaiman (2022).
+        *Physics-informed attention temporal convolutional network for EEG-based motor imagery classification.*
+        IEEE Transactions on Industrial Informatics. doi:10.1109/TII.2022.3197419.
+    .. [2] Official EEG-ATCNet implementation (TensorFlow):
+        https://github.com/Altaheri/EEG-ATCNet/blob/main/models.py
     """
 
     def __init__(
@@ -231,7 +369,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
                 nn.Sequential(
                     *[
                         _TCNResidualBlock(
-                            in_channels=self.F2,
+                            in_channels=self.F2 if i == 0 else self.tcn_n_filters,
                             kernel_size=self.tcn_kernel_size,
                             n_filters=self.tcn_n_filters,
                             dropout=self.tcn_dropout,
@@ -249,7 +387,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
             self.final_layer = nn.ModuleList(
                 [
                     MaxNormLinear(
-                        in_features=self.F2 * self.n_windows,
+                        in_features=self.tcn_n_filters * self.n_windows,
                         out_features=self.n_outputs,
                         max_norm_val=self.max_norm_const,
                     )
@@ -259,7 +397,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
             self.final_layer = nn.ModuleList(
                 [
                     MaxNormLinear(
-                        in_features=self.F2,
+                        in_features=self.tcn_n_filters,
                         out_features=self.n_outputs,
                         max_norm_val=self.max_norm_const,
                     )
@@ -269,7 +407,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
 
         self.out_fun = nn.Identity()
 
-    def forward(self, X):
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
         # Dimension: (batch_size, C, T)
         X = self.ensuredims(X)
         # Dimension: (batch_size, C, T, 1)
@@ -556,7 +694,8 @@ class _TCNResidualBlock(nn.Module):
         # Reshape the input for the residual connection when necessary
         if in_channels != n_filters:
             self.reshaping_conv = nn.Conv1d(
-                n_filters,
+                in_channels=in_channels,  # Specify input channels
+                out_channels=n_filters,  # Specify output channels
                 kernel_size=1,
                 padding="same",
             )
@@ -576,7 +715,7 @@ class _TCNResidualBlock(nn.Module):
         out = self.activation(out)
         out = self.drop2(out)
 
-        out = self.reshaping_conv(out)
+        X = self.reshaping_conv(X)
 
         # ----- Residual connection -----
         out = X + out