braindecode 1.8.0.dev174836782__tar.gz → 1.8.0.dev180935432__tar.gz

{braindecode-1.8.0.dev174836782/braindecode.egg-info → braindecode-1.8.0.dev180935432}/PKG-INFO

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

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/augmentation/base.py

@@ -3,6 +3,7 @@
 #          Bruno Aristimunha <b.aristimunha@gmail.com>
 #          Martin Wimpff <martin.wimpff@iss.uni-stuttgart.de>
 #          Valentin Iovene <val@too.gy>
+#          Sarthak Tayal <sarthaktayal2@gmail.com>
 # License: BSD (3-clause)
 
 from numbers import Real
@@ -176,6 +177,14 @@ class Compose(Transform):
         return X, y
 
 
+def _as_mixed_target(y, lam_dtype):
+    # an untouched target is the same as being mixed with itself with lam of one
+    if isinstance(y, (tuple, list)) and len(y) == 3:
+        return tuple(y)
+    lam = torch.ones(y.shape[0], device=y.device, dtype=lam_dtype)
+    return y, y, lam
+
+
 class _AugmentationCollate:
     """Collate that applies a transform to each batch, with optional expansion.
 
@@ -193,7 +202,10 @@ class _AugmentationCollate:
         ``0`` (default) applies the transform in place (batch size unchanged).
         ``> 0`` keeps the clean originals and appends ``n_augmentation``
         independently transformed copies, returning ``(X, y)`` of
-        ``(1 + n_augmentation)`` times the original size.
+        ``(1 + n_augmentation)`` times the original size. When the transform
+        mixes targets, as :class:`braindecode.augmentation.Mixup` does, ``y``
+        stays the ``(y_a, y_b, lam)`` triple and the clean originals get a
+        mixing coefficient of one.
     """
 
     def __init__(self, transform, device=None, n_augmentation=0):
@@ -216,6 +228,18 @@ class _AugmentationCollate:
             aug_X, aug_y = self.transform(X, y)
             xs.append(aug_X)
             ys.append(aug_y)
+        mixed_ys = [
+            aug_y
+            for aug_y in ys
+            if isinstance(aug_y, (tuple, list)) and len(aug_y) == 3
+        ]
+        if mixed_ys:
+            # a target-mixing transform such as Mixup returns (y_a, y_b, lam),
+            # so the parts are concatenated one by one and the untouched copies
+            # are given a mixing coefficient of one
+            lam_dtype = mixed_ys[0][2].dtype
+            ys = [_as_mixed_target(aug_y, lam_dtype) for aug_y in ys]
+            return torch.cat(xs), tuple(torch.cat(part) for part in zip(*ys))
         return torch.cat(xs), torch.cat(ys)
 
 

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/augmentation/functional.py

@@ -1065,13 +1065,11 @@ def mixup(
     batch_size, n_channels, n_times = X.shape
 
     X_mix = torch.zeros((batch_size, n_channels, n_times)).to(device)
-    y_a = torch.arange(batch_size).to(device)
-    y_b = torch.arange(batch_size).to(device)
+    y_a = y.clone()
+    y_b = y[idx_perm].clone()
 
     for idx in range(batch_size):
         X_mix[idx] = lam[idx] * X[idx] + (1 - lam[idx]) * X[idx_perm[idx]]
-        y_a[idx] = y[idx]
-        y_b[idx] = y[idx_perm[idx]]
 
     return X_mix, (y_a, y_b, lam)
 

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/augmentation/transforms.py

@@ -1068,16 +1068,18 @@ class Mixup(Transform):
         device = X.device
         batch_size, _, _ = X.shape
 
+        # lam follows the dtype of X, numpy draws float64 and that would leak
+        # into the mixed signal and into the loss returned by mixup_criterion
         if self.alpha > 0:
             if self.beta_per_sample:
                 lam = torch.as_tensor(
                     self.rng.beta(self.alpha, self.alpha, batch_size)
-                ).to(device)
+                ).to(device=device, dtype=X.dtype)
             else:
-                lam = torch.ones(batch_size).to(device)
+                lam = torch.ones(batch_size, dtype=X.dtype).to(device)
                 lam *= self.rng.beta(self.alpha, self.alpha)
         else:
-            lam = torch.ones(batch_size).to(device)
+            lam = torch.ones(batch_size, dtype=X.dtype).to(device)
 
         idx_perm = torch.as_tensor(
             self.rng.permutation(

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/datasets/base.py

@@ -59,6 +59,7 @@ def _html_row(label, value):
 
 _METADATA_INTERNAL_COLS = {
     "i_window_in_trial",
+    "i_trial_in_dataset",
     "i_start_in_trial",
     "i_stop_in_trial",
     "target",
@@ -1374,6 +1375,20 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
                 "datasets are WindowsDataset."
             )
 
+        for ds in self.datasets:
+            if hasattr(ds, "_windows") and ds._windows is not None:
+                df = ds._windows.metadata
+            else:
+                df = ds.metadata
+            if (
+                "i_trial_in_dataset" in df.columns
+                and "i_trial_in_dataset" in ds.description
+            ):
+                raise ValueError(
+                    "Dataset descriptions cannot contain the reserved window "
+                    "metadata key 'i_trial_in_dataset'."
+                )
+
         all_dfs = list()
         for ds in self.datasets:
             if hasattr(ds, "_windows") and ds._windows is not None:

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/datasets/bids/hub.py

@@ -28,6 +28,7 @@ The format follows a BIDS-inspired sourcedata structure:
 # License: BSD (3-clause)
 
 import contextlib
+import copy
 import json
 import logging
 import tempfile
@@ -56,6 +57,7 @@ from .hub_io import (
     _load_eegwindows_from_zarr,
     _load_raw_from_zarr,
     _load_windows_from_zarr,
+    _prepare_info_for_json,
     _save_eegwindows_to_zarr,
     _save_raw_to_zarr,
     _save_windows_to_zarr,
@@ -72,6 +74,33 @@ log = logging.getLogger(__name__)
 _LOCK_FILE = "format_info.json"
 
 
+def _normalize_kwargs_for_json(value, field_name):
+    """Return one preprocessing-kwargs value as native strict JSON."""
+
+    def _convert(obj):
+        if isinstance(obj, np.ndarray):
+            return _convert(obj.tolist())
+        if isinstance(obj, dict):
+            return {key: _convert(item) for key, item in obj.items()}
+        if isinstance(obj, (list, tuple)):
+            return [_convert(item) for item in obj]
+        if isinstance(obj, np.generic):
+            return _convert(obj.item())
+        return obj
+
+    try:
+        converted = _convert(value)
+        if isinstance(converted, str):
+            raise ValueError(
+                "a root string is ambiguous with the legacy encoded format"
+            )
+        return json.loads(json.dumps(converted, allow_nan=False))
+    except (OverflowError, RecursionError, TypeError, ValueError) as error:
+        raise ValueError(
+            f"{field_name} must contain only finite JSON-serializable values"
+        ) from error
+
+
 class HubDatasetMixin:
     """
     Mixin class for Hugging Face Hub integration with EEG datasets.
@@ -685,14 +714,57 @@ class HubDatasetMixin:
                 f"{output_path} already exists. Set overwrite=True to replace it."
             )
 
-        # Create zarr store (zarr v3 API)
-        root = zarr.open(str(output_path), mode="w")
-
         # Validate uniformity across all datasets using shared validation
         dataset_type, _, _ = hub_validation.validate_dataset_uniformity(self.datasets)
 
-        # Keep reference to first dataset for preprocessing kwargs
-        first_ds = self.datasets[0]
+        # Normalize every JSON value before opening the output store. The cached
+        # values below are the exact values handed to the write helpers.
+        prepared_infos = []
+        for ds in self.datasets:
+            if dataset_type == "WindowsDataset":
+                info = ds.windows.info
+            elif dataset_type in ("EEGWindowsDataset", "RawDataset"):
+                info = ds.raw.info
+            prepared_infos.append(_prepare_info_for_json(info.to_json_dict()))
+
+        prepared_kwargs = {}
+        for kwarg_name in [
+            "raw_preproc_kwargs",
+            "window_kwargs",
+            "window_preproc_kwargs",
+        ]:
+            expected_present = hasattr(self.datasets[0], kwarg_name)
+            expected_value = None
+            expected_token = None
+            for i_ds, ds in enumerate(self.datasets):
+                present = hasattr(ds, kwarg_name)
+                if present != expected_present:
+                    raise ValueError(
+                        f"{kwarg_name} on dataset {i_ds} has inconsistent presence; "
+                        "the Zarr format stores one global value"
+                    )
+                if not present:
+                    continue
+                value = _normalize_kwargs_for_json(
+                    getattr(ds, kwarg_name), f"{kwarg_name} on dataset {i_ds}"
+                )
+                value_token = json.dumps(
+                    value, sort_keys=True, separators=(",", ":"), allow_nan=False
+                )
+                if i_ds == 0:
+                    expected_value = value
+                    expected_token = value_token
+                elif value_token != expected_token:
+                    raise ValueError(
+                        f"{kwarg_name} on dataset {i_ds} differs from dataset 0; "
+                        "the Zarr format stores one global value"
+                    )
+            if expected_present:
+                prepared_kwargs[kwarg_name] = expected_value
+
+        # Create compressor and zarr store (zarr v3 API) only after preflight.
+        compressor = _create_compressor(compression, compression_level)
+        root = zarr.open(str(output_path), mode="w")
 
         # Store global metadata
         root.attrs["n_datasets"] = len(self.datasets)
@@ -706,21 +778,10 @@ class HubDatasetMixin:
         root.attrs["zarr_version"] = zarr.__version__
         root.attrs["scipy_version"] = scipy.__version__
 
-        # Save preprocessing kwargs (check first dataset, assuming uniform preprocessing)
-        # These are typically set by windowing functions on individual datasets
-        for kwarg_name in [
-            "raw_preproc_kwargs",
-            "window_kwargs",
-            "window_preproc_kwargs",
-        ]:
-            # Check first dataset for these attributes
-            if hasattr(first_ds, kwarg_name):
-                kwargs = getattr(first_ds, kwarg_name)
-                if kwargs:
-                    root.attrs[kwarg_name] = json.dumps(kwargs)
-
-        # Create compressor
-        compressor = _create_compressor(compression, compression_level)
+        # Save preprocessing kwargs from the preflight cache. These are
+        # typically set by windowing functions on individual datasets.
+        for kwarg_name, kwargs in prepared_kwargs.items():
+            root.attrs[kwarg_name] = kwargs
 
         # Save each recording
         for i_ds, ds in enumerate(self.datasets):
@@ -731,7 +792,7 @@ class HubDatasetMixin:
                 data = ds.windows.get_data()
                 metadata = ds.windows.metadata
                 description = ds.description
-                info_dict = ds.windows.info.to_json_dict()
+                info_dict = prepared_infos[i_ds]
                 target_name = ds.target_name if hasattr(ds, "target_name") else None
 
                 # Save using inlined function
@@ -750,7 +811,7 @@ class HubDatasetMixin:
                 raw = ds.raw
                 metadata = ds.metadata
                 description = ds.description
-                info_dict = ds.raw.info.to_json_dict()
+                info_dict = prepared_infos[i_ds]
                 targets_from = ds.targets_from
                 last_target_only = ds.last_target_only
 
@@ -771,7 +832,7 @@ class HubDatasetMixin:
                 # Get continuous raw data from RawDataset
                 raw = ds.raw
                 description = ds.description
-                info_dict = ds.raw.info.to_json_dict()
+                info_dict = prepared_infos[i_ds]
                 target_name = ds.target_name if hasattr(ds, "target_name") else None
 
                 # Save using inlined function
@@ -956,10 +1017,14 @@ class HubDatasetMixin:
             "window_preproc_kwargs",
         ]:
             if kwarg_name in root.attrs:
-                kwargs = json.loads(root.attrs[kwarg_name])
+                kwargs = root.attrs[kwarg_name]
+                if isinstance(kwargs, str):
+                    # Stores written by older braindecode versions kept
+                    # these attributes as double-encoded JSON strings.
+                    kwargs = json.loads(kwargs)
                 # Set on each individual dataset (where they were originally stored)
                 for ds in datasets:
-                    setattr(ds, kwarg_name, kwargs)
+                    setattr(ds, kwarg_name, copy.deepcopy(kwargs))
 
         return concat_ds
 

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/datasets/bids/hub_io.py

@@ -8,6 +8,7 @@ These functions keep the Zarr serialization details isolated from hub.py.
 from __future__ import annotations
 
 import json
+from numbers import Real
 from pathlib import Path
 
 import numpy as np
@@ -17,17 +18,78 @@ from mne.utils import _soft_import
 zarr = _soft_import("zarr", purpose="hugging face integration", strict=False)
 
 
+def _is_non_bool_real(value):
+    return not isinstance(value, (bool, np.bool_)) and isinstance(
+        value, (Real, np.integer, np.floating)
+    )
+
+
+def _prepare_info_for_json(obj, path="info", *, _in_numeric_sequence=False):
+    """Normalize an MNE Info value to strict JSON without losing sequence NaNs."""
+    if isinstance(obj, np.ndarray):
+        obj = obj.tolist()
+
+    if isinstance(obj, dict):
+        return {
+            key: _prepare_info_for_json(value, f"{path}.{key}")
+            for key, value in obj.items()
+        }
+
+    if isinstance(obj, (list, tuple)):
+        values = list(obj)
+        has_none = any(value is None for value in values)
+        has_number = any(_is_non_bool_real(value) for value in values)
+        all_none = bool(values) and all(value is None for value in values)
+        if all_none or (has_none and has_number):
+            raise ValueError(
+                f"{path} is ambiguous: numeric sequences cannot contain JSON null"
+            )
+
+        is_numeric_sequence = bool(values) and all(
+            _is_non_bool_real(value) for value in values
+        )
+        return [
+            _prepare_info_for_json(
+                value,
+                f"{path}[{index}]",
+                _in_numeric_sequence=is_numeric_sequence,
+            )
+            for index, value in enumerate(values)
+        ]
+
+    if isinstance(obj, (bool, np.bool_)):
+        return bool(obj)
+    if isinstance(obj, (int, np.integer)):
+        return int(obj)
+    if _is_non_bool_real(obj):
+        value = float(obj)
+        if np.isnan(value):
+            if _in_numeric_sequence:
+                return None
+            raise ValueError(f"{path} contains unsupported NaN")
+        if np.isposinf(value):
+            raise ValueError(f"{path} contains positive infinity")
+        if np.isneginf(value):
+            raise ValueError(f"{path} contains negative infinity")
+        return value
+    if obj is None or isinstance(obj, str):
+        return obj
+    raise ValueError(f"{path} contains a non-JSON-serializable value")
+
+
 def _restore_nan_from_json(obj):
-    """Restore NaN values from None in legacy zarr stores.
+    """Restore NaN values from None in JSON-loaded attributes.
 
-    Datasets saved before zarr v3 native NaN support used
-    ``_sanitize_for_json`` to convert NaN/Inf → None. This restores them
-    on load so ``mne.Info.from_json_dict`` gets proper NaN arrays.
+    JSON null is reserved for NaN only in non-boolean numeric sequences.
+    All-null lists are therefore the representation of validated all-NaN
+    sequences, while ordinary null values elsewhere remain unchanged.
     """
     if isinstance(obj, dict):
         return {k: _restore_nan_from_json(v) for k, v in obj.items()}
     if isinstance(obj, list):
-        if len(obj) > 0 and all(isinstance(x, (int, float, type(None))) for x in obj):
+        if len(obj) > 0 and all(
+            value is None or _is_non_bool_real(value) for value in obj
+        ):
             return [np.nan if x is None else x for x in obj]
         return [_restore_nan_from_json(v) for v in obj]
     return obj

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/datasets/sleep_physionet.py

@@ -117,9 +117,15 @@ class SleepPhysionet(BaseConcatDataset):
             sleep_event_inds = np.where(mask)[0]
 
             # Crop raw
-            tmin = annots[int(sleep_event_inds[0])]["onset"] - crop_wake_mins * 60
-            tmax = annots[int(sleep_event_inds[-1])]["onset"] + crop_wake_mins * 60
-            raw.crop(tmin=max(tmin, raw.times[0]), tmax=min(tmax, raw.times[-1]))
+            a_tmin = annots[sleep_event_inds[0]]
+            a_tmax = annots[sleep_event_inds[-1]]
+            tmin = a_tmin["onset"] - crop_wake_mins * 60
+            tmax = a_tmax["onset"] + a_tmax["duration"] + crop_wake_mins * 60
+            raw.crop(
+                tmin=max(tmin, raw.times[0]),
+                tmax=min(tmax, raw.times[-1] + 1 / raw.info["sfreq"]),
+                include_tmax=False,
+            )
 
         # Rename EEG channels
         ch_names = {i: i.replace("EEG ", "") for i in raw.ch_names if "EEG" in i}

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/models/atcnet.py

@@ -196,7 +196,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
     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
+    att_drop_prob : float
         Dropout probability used in the attention block, denoted pa in table 1
         of the paper [1]_. Defaults to 0.5 as in [1]_.
     tcn_depth : int
@@ -206,7 +206,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
     tcn_kernel_size : int
         Temporal kernel size used in TCN block, denoted Kt in table 1 of the
         paper [1]_. Defaults to 4 as in [1]_.
-    tcn_dropout : float
+    tcn_drop_prob : float
         Dropout probability used in the TCN block, denoted pt in table 1
         of the paper [1]_. Defaults to 0.3 as in [1]_.
     tcn_activation : torch.nn.Module

{braindecode-1.8.0.dev174836782 → braindecode-1.8.0.dev180935432}/braindecode/models/attn_sleep.py

@@ -1,10 +1,36 @@
 # Authors: Divyesh Narayanan <divyesh.narayanan@gmail.com>
+#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
+#
+# This implementation derives from https://github.com/emadeldeen24/AttnSleep:
+#
+# MIT License
+#
+# Copyright (c) 2020 Emadeldeen Eldele
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
 
 import math
 import warnings
 from copy import deepcopy
+from numbers import Integral
 
 import torch
 import torch.nn.functional as F
@@ -33,7 +59,8 @@ class AttnSleep(EEGModuleMixin, nn.Module):
 
     Warning - This model was designed for signals of 30 seconds at 100Hz or 125Hz (in which case
     the reference architecture from [1]_ which was validated on SHHS dataset [2]_ will be used)
-    to use any other input is likely to make the model perform in unintended ways.
+    to use any other input is likely to make the model perform in unintended ways. Any other
+    window length also needs a ``d_model`` of its own, see the parameter below.
 
     Parameters
     ----------
@@ -44,7 +71,10 @@ class AttnSleep(EEGModuleMixin, nn.Module):
         Also the input dimension of the first FC layer in the feed forward
         and the output of the second FC layer in the same.
         Increase for higher sampling rate/signal length.
-        It should be divisible by n_attn_heads
+        It should be divisible by n_attn_heads. It must also equal the number
+        of time steps returned by the feature extractor: 80 for 30 seconds at
+        100 Hz and 100 for 30 seconds at 125 Hz. A construction error reports
+        the value needed for other window lengths.
     d_ff : int
         Output dimension of the first FC layer in the feed forward and the
         input dimension of the second FC layer in the same.
@@ -58,15 +88,12 @@ class AttnSleep(EEGModuleMixin, nn.Module):
         If True, return the features, i.e. the output of the feature extractor
         (before the final linear layer). If False, pass the features through
         the final linear layer.
-    n_classes : int
-        Alias for `n_outputs`.
-    input_size_s : float
-        Alias for `input_window_seconds`.
     activation : nn.Module, default=nn.ReLU
-        Activation function class to apply. Should be a PyTorch activation
+        Activation function class to apply in the AFR block and the TCE
+        feed-forward block. Should be a PyTorch activation
         module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
-    activation_mrcnn : nn.Module, default=nn.ReLU
-        Activation function class to apply in the Mask R-CNN layer.
+    activation_mrcnn : nn.Module, default=nn.GELU
+        Activation function class to apply in the multi-resolution CNN layer.
         Should be a PyTorch activation module class like ``nn.ReLU`` or
         ``nn.GELU``. Default is ``nn.GELU``.
 
@@ -99,6 +126,14 @@ class AttnSleep(EEGModuleMixin, nn.Module):
         n_chans=None,
         n_times=None,
     ):
+        if (
+            sum(value is not None for value in (n_times, sfreq, input_window_seconds))
+            < 2
+        ):
+            raise ValueError(
+                "AttnSleep requires at least two of n_times, sfreq, and "
+                "input_window_seconds."
+            )
         super().__init__(
             n_outputs=n_outputs,
             n_chans=n_chans,
@@ -140,6 +175,14 @@ class AttnSleep(EEGModuleMixin, nn.Module):
             activation=activation_mrcnn,
             activation_se=activation,
         )
+        feature_length = self._feature_length(mrcnn, self.n_times)
+        if feature_length != d_model:
+            raise ValueError(
+                f"d_model is {d_model} but the feature extractor returns "
+                f"{feature_length} time steps for an input of {self.n_times} "
+                f"samples at {self.sfreq} Hz. Set d_model={feature_length}, with "
+                "an n_attn_heads that divides it."
+            )
         attn = _MultiHeadedAttention(n_attn_heads, d_model, after_reduced_cnn_size)
         ff = _PositionwiseFeedForward(d_model, d_ff, drop_prob, activation=activation)
         tce = _TCE(
@@ -150,23 +193,26 @@ class AttnSleep(EEGModuleMixin, nn.Module):
         )
 
         self.feature_extractor = nn.Sequential(mrcnn, tce)
-        self.len_last_layer = self._len_last_layer(self.n_times)
+        self.len_last_layer = feature_length * after_reduced_cnn_size
         self.return_feats = return_feats
 
         # TODO: Add new way to handle return features
         """if return_feats:
             raise ValueError("return_feat == True is not accepted anymore")"""
         if not return_feats:
-            self.final_layer = nn.Linear(
-                d_model * after_reduced_cnn_size, self.n_outputs
-            )
-
-    def _len_last_layer(self, input_size):
-        self.feature_extractor.eval()
-        with torch.no_grad():
-            out = self.feature_extractor(torch.Tensor(1, 1, input_size))
-        self.feature_extractor.train()
-        return len(out.flatten())
+            self.final_layer = nn.Linear(self.len_last_layer, self.n_outputs)
+
+    @staticmethod
+    def _feature_length(mrcnn, n_times):
+        training_states = [(module, module.training) for module in mrcnn.modules()]
+        mrcnn.eval()
+        try:
+            with torch.no_grad():
+                out = mrcnn(torch.zeros(1, 1, n_times))
+        finally:
+            for module, was_training in training_states:
+                module.training = was_training
+        return out.shape[-1]
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
@@ -195,7 +241,7 @@ class _SELayer(nn.Module):
         self.avg_pool = nn.AdaptiveAvgPool1d(1)
         self.fc = nn.Sequential(
             nn.Linear(channel, channel // reduction, bias=False),
-            activation(inplace=True),
+            activation(),
             nn.Linear(channel // reduction, channel, bias=False),
             nn.Sigmoid(),
         )
@@ -236,10 +282,10 @@ class _SEBasicBlock(nn.Module):
         super(_SEBasicBlock, self).__init__()
         self.conv1 = nn.Conv1d(inplanes, planes, stride)
         self.bn1 = nn.BatchNorm1d(planes)
-        self.relu = activation(inplace=True)
+        self.relu = activation()
         self.conv2 = nn.Conv1d(planes, planes, 1)
         self.bn2 = nn.BatchNorm1d(planes)
-        self.se = _SELayer(planes, reduction)
+        self.se = _SELayer(planes, reduction, activation=activation)
         self.downsample = downsample
         self.stride = stride
         self.features = nn.Sequential(
@@ -320,11 +366,11 @@ class _MRCNN(nn.Module):
         self.dropout = nn.Dropout(drate)
         self.inplanes = 128
         self.AFR = self._make_layer(
-            _SEBasicBlock, after_reduced_cnn_size, 1, activate=activation_se
+            _SEBasicBlock, after_reduced_cnn_size, 1, activation=activation_se
         )
 
     def _make_layer(
-        self, block, planes, blocks, stride=1, activate: type[nn.Module] = nn.ReLU
+        self, block, planes, blocks, stride=1, activation: type[nn.Module] = nn.ReLU
     ):  # makes residual SE block
         downsample = None
         if stride != 1 or self.inplanes != planes * block.expansion:
@@ -340,10 +386,12 @@ class _MRCNN(nn.Module):
             )
 
         layers = []
-        layers.append(block(self.inplanes, planes, stride, downsample))
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, activation=activation)
+        )
         self.inplanes = planes * block.expansion
         for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes, activate=activate))
+            layers.append(block(self.inplanes, planes, activation=activation))
 
         return nn.Sequential(*layers)
 
@@ -375,9 +423,20 @@ class _MultiHeadedAttention(nn.Module):
     def __init__(self, h, d_model, after_reduced_cnn_size, dropout=0.1):
         """Take in model size and number of heads."""
         super().__init__()
-        assert d_model % h == 0
+        if (
+            isinstance(h, bool)
+            or not isinstance(h, Integral)
+            or h <= 0
+            or d_model % h != 0
+        ):
+            raise ValueError(
+                "n_attn_heads must be a positive integer that divides d_model, "
+                f"got n_attn_heads={h!r} and d_model={d_model}."
+            )
+        h = int(h)
         self.d_per_head = d_model // h
         self.h = h
+        self.attn = torch.empty(0)
 
         base_conv = CausalConv1d(
             in_channels=after_reduced_cnn_size,