braindecode 1.8.0.dev1128__tar.gz → 1.8.0.dev169056309__tar.gz

braindecode-1.8.0.dev169056309/NOTICE.txt

@@ -0,0 +1,65 @@
+# BRAINDECODE Notice
+
+## Licensed Components
+
+### BSD-3-Clause Licensed Files
+
+All files within the `braindecode/` package are licensed under the BSD-3-Clause
+License, except for those listed in the sections below.
+
+### CC BY-NC 4.0 Licensed Files
+
+The following components are licensed under the Creative Commons Attribution-NonCommercial 4.0 International License:
+
+- `braindecode/models/eegminer.py`
+- `braindecode/models/meta_neuromotor.py`
+- `braindecode/models/brainmodule.py`
+
+As well as class later imported into the `braindecode.models.module` named as GeneralizedGaussianFilter.
+
+The `meta_neuromotor.py` file is a derivative of
+`facebookresearch/generic-neuromotor-interface`, released by Meta Platforms,
+Inc. under CC BY-NC 4.0, and inherits the same noncommercial terms.
+
+### CC BY-NC-SA 4.0 Licensed Files
+
+The following components are licensed under the Creative Commons
+Attribution-NonCommercial-ShareAlike 4.0 International License:
+
+- `braindecode/models/emg2qwerty.py`
+
+The `emg2qwerty.py` file is a derivative of `facebookresearch/emg2qwerty`,
+released by Meta Platforms, Inc. under CC BY-NC-SA 4.0, and inherits the
+same noncommercial ShareAlike terms.
+
+### MIT Licensed Files
+
+The following components are licensed under the MIT License:
+
+- `braindecode/models/ctnet.py`
+- `braindecode/models/dance.py`
+- `braindecode/models/medformer.py`
+- `braindecode/models/tcformer.py`
+- `braindecode/models/ifnet.py`
+
+### Apache-2.0 Licensed Files
+
+The following components are licensed under the Apache License 2.0:
+
+- `braindecode/models/mvpformer.py`
+- `braindecode/models/zuna.py`
+- `braindecode/models/luna.py`
+
+## License Links
+
+- [BSD-3-Clause License](https://opensource.org/licenses/BSD-3-Clause)
+- [CC BY-NC 4.0 License](https://creativecommons.org/licenses/by-nc/4.0/)
+- [CC BY-NC-SA 4.0 License](https://creativecommons.org/licenses/by-nc-sa/4.0/)
+- [MIT License](https://opensource.org/licenses/MIT)
+- [Apache-2.0 License](https://www.apache.org/licenses/LICENSE-2.0)
+
+## Note
+
+This list covers files whose own headers declare a license other than
+BSD-3-Clause. Per-file provenance review of the remaining models is
+tracked separately.

{braindecode-1.8.0.dev1128/braindecode.egg-info → braindecode-1.8.0.dev169056309}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.0.dev1128
+Version: 1.8.0.dev169056309
 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>
@@ -34,7 +34,6 @@ Requires-Dist: pandas
 Requires-Dist: wfdb>=4.3.1
 Requires-Dist: linear_attention_transformer
 Requires-Dist: docstring_inheritance
-Requires-Dist: rotary_embedding_torch
 Requires-Dist: pydantic>=2.0
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.4.3; extra == "moabb"
@@ -51,6 +50,7 @@ Requires-Dist: pytest-cov; extra == "tests"
 Requires-Dist: codecov; extra == "tests"
 Requires-Dist: pytest_cases; extra == "tests"
 Requires-Dist: mypy; extra == "tests"
+Requires-Dist: ipython; extra == "tests"
 Requires-Dist: transformers>=4.57.0; extra == "tests"
 Requires-Dist: bids_validator; extra == "tests"
 Provides-Extra: typing
@@ -260,7 +260,10 @@ This project is primarily licensed under the BSD-3-Clause License.
 Additional Components
 =====================
 
-Some components within this repository are licensed under the Creative Commons
-Attribution-NonCommercial 4.0 International License.
+Some components within this repository are licensed under other licenses, including
+Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0), Creative
+Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0), MIT
+and Apache-2.0.
 
-Please refer to the ``LICENSE`` and ``NOTICE`` files for more detailed information.
+Please refer to the ``LICENSE`` and ``NOTICE`` files for the per-file list and more
+detailed information.

{braindecode-1.8.0.dev1128 → braindecode-1.8.0.dev169056309}/README.rst

@@ -170,7 +170,10 @@ This project is primarily licensed under the BSD-3-Clause License.
 Additional Components
 =====================
 
-Some components within this repository are licensed under the Creative Commons
-Attribution-NonCommercial 4.0 International License.
+Some components within this repository are licensed under other licenses, including
+Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0), Creative
+Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0), MIT
+and Apache-2.0.
 
-Please refer to the ``LICENSE`` and ``NOTICE`` files for more detailed information.
+Please refer to the ``LICENSE`` and ``NOTICE`` files for the per-file list and more
+detailed information.

braindecode-1.8.0.dev169056309/braindecode/datasets/_notebook_viewer.py

@@ -0,0 +1,255 @@
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# License: BSD (3-clause)
+"""Serverless in-notebook viewer for file-backed recordings.
+
+The bytes on disk behind a dataset element are inlined in the cell output
+as base64 and handed to the deployed eegdash-viewer over its ``postMessage``
+bridge (``docs/embedding.md`` in https://github.com/eegdash/eegdash-viewer):
+no server, no CORS. The output is an iframe plus an inline script, so it
+renders when the cell ran in your session or the saved notebook is trusted
+(``jupyter trust``).
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import uuid
+from pathlib import Path
+from urllib.parse import urlsplit
+
+import mne_bids
+from mne.utils import _soft_import
+
+CDN = "https://eegdash.github.io/eegdash-viewer"
+MAX_BYTES = 64 * 2**20  # base64 output per call; it is saved with the notebook
+EXTENSIONS = {
+    ".set",
+    ".edf",
+    ".bdf",
+    ".vhdr",
+    ".fif",
+    ".snirf",
+    ".nwb",
+}  # viewer readers
+_BIDS = {"eeg", "ieeg", "emg", "meg", "nirs"}
+_SIBLINGS = {".vhdr": (".eeg", ".vmrk"), ".set": (".fdt",)}  # travel with the header
+_HEADER = {s: h for h, ss in _SIBLINGS.items() for s in ss}  # data file -> header
+
+_SCRIPT = """<iframe id=%(id)s title="eegdash trace viewer" style="width:100%%;height:%(height)spx;
+border:1px solid var(--jp-border-color1,#d9dce1);border-radius:6px;background:transparent"></iframe>
+<script>
+(function () {
+  var self = document.currentScript, id = %(id)s;   // Lab re-runs scripts in place; VS Code/nbclassic elsewhere
+  var frame = (self && self.previousElementSibling && self.previousElementSibling.tagName === "IFRAME")
+    ? self.previousElementSibling : document.getElementById(id);
+  if (!frame) { console.error("eegdash viewer: output iframe " + id + " not found"); return; }
+  var payload = %(payload)s, origin = %(origin)s, files = null, pose = null;
+  function decode(b64) {
+    if (Uint8Array.fromBase64) return Uint8Array.fromBase64(b64);
+    var bin = atob(b64), out = new Uint8Array(bin.length);
+    for (var i = 0; i < bin.length; i++) out[i] = bin.charCodeAt(i);
+    return out;
+  }
+  function send(target) {
+    try {
+      if (!files) {
+        files = payload.files.map(function (f) { return new File([decode(f.b64)], f.name); });
+        pose = payload.pose ? "data:application/json;base64," + payload.pose : null;
+        payload = null;
+      }
+      frame.contentWindow.postMessage({ type: "eegdash-viewer:open", files: files, pose: pose }, target || origin);
+    } catch (err) {
+      frame.insertAdjacentHTML("afterend", '<div style="font:12px system-ui;color:#b3261e">eegdash viewer: '
+        + String(err.message).replace(/</g, "&lt;") + "</div>");
+    }
+  }
+  window.addEventListener("message", function onMessage(e) {
+    if (e.source === frame.contentWindow && e.data && e.data.type === "eegdash-viewer:ready") send(e.origin);
+    else if (!frame.isConnected) window.removeEventListener("message", onMessage);
+  });
+  frame.src = %(src)s;   // after the listener, so "ready" can never precede it
+})();
+</script>"""
+
+
+def recording_files(recording: Path) -> tuple[list[Path], Path | None]:
+    """``(files, pose)`` to inline for a recording: the header first, then the
+    split-format siblings, the BIDS-inherited ``_channels.tsv``/``_events.tsv``
+    (``mne_bids`` parses the name; plain names get none) and, separately, the
+    ``<prefix>_desc-pose.json`` hand-pose sidecar next to it. Symlinks keep
+    their name (git-annex/datalad)."""
+    rec = Path(recording)
+    if not rec.exists():
+        raise ValueError(
+            f"{rec.name}: file not found (a datalad symlink may need `datalad get`)"
+        )
+    if (
+        rec.is_dir()
+        or rec.suffix.lower() not in EXTENSIONS
+        or rec.stem.endswith("_epo")
+    ):
+        raise ValueError(
+            f"{rec.name}: the viewer opens raw recordings in {' '.join(sorted(EXTENSIONS))}"
+        )
+    sidecars: list[Path | None] = []
+    try:
+        bids = mne_bids.get_bids_path_from_fname(rec, check=False)
+        if bids.subject is not None:  # hyphen-free names parse, but are not BIDS
+            sidecars = [
+                bids.find_matching_sidecar(
+                    suffix=s, extension=".tsv", on_error="ignore"
+                )
+                for s in ("channels", "events")
+            ]
+    except (
+        KeyError,
+        ValueError,
+    ):  # not a BIDS name / unknown entity somewhere in the tree
+        pass
+    files = [rec]
+    for p in [
+        rec.with_suffix(e) for e in _SIBLINGS.get(rec.suffix.lower(), ())
+    ] + sidecars:
+        if p and p.is_file():
+            if p not in files:
+                files.append(p)
+        elif p and p.is_symlink():  # dangling: git-annex/datalad content not fetched
+            raise ValueError(f"{p.name}: dangling symlink (try `datalad get`)")
+    stem, _, token = rec.stem.rpartition("_")
+    pose = rec.with_name((stem if token in _BIDS else rec.stem) + "_desc-pose.json")
+    return files, pose if pose.is_file() else None
+
+
+def build_viewer_html(
+    recording: Path,
+    *,
+    height: int = 520,
+    cdn_url: str = CDN,
+    max_bytes: int = MAX_BYTES,
+) -> str:
+    """Viewer iframe + inlined bytes + bridge script for one recording."""
+    url = urlsplit(cdn_url)
+    if (
+        url.scheme not in ("http", "https")
+        or not url.netloc
+        or url.username
+        or url.query
+        or url.fragment
+        or url.path.endswith(("index.html", "index.htm"))
+    ):
+        raise ValueError(
+            f"cdn_url must be the viewer's base http(s) URL, got {cdn_url!r}"
+        )
+    files, pose = recording_files(recording)
+    encoded = sum(
+        4 * -(-p.stat().st_size // 3) for p in files + ([pose] if pose else [])
+    )
+    if encoded > max_bytes:
+        raise ValueError(
+            f"{files[0].name}: {encoded / 2**20:.1f} MiB of base64 would be inlined into the "
+            f"notebook output (max_bytes={max_bytes / 2**20:.1f} MiB); crop/downsample or raise it"
+        )
+    rec = files[0]
+    # The viewer picks the recording by its *_<datatype>.<ext> name: plain names
+    # are posted as <stem>_eeg<ext>, an EEGLAB .fdt next to the posted .set name.
+    head = (
+        rec.name
+        if rec.stem.rpartition("_")[2] in _BIDS
+        else f"{rec.stem}_eeg{rec.suffix.lower()}"
+    )
+    names = [head] + [
+        head.rsplit("_", 1)[0] + "_eeg.fdt" if p.suffix.lower() == ".fdt" else p.name
+        for p in files[1:]
+    ]
+    b64 = [base64.b64encode(p.read_bytes()).decode() for p in files]
+    literals = {
+        "id": f"eegdash-viewer-{uuid.uuid4().hex[:8]}",
+        "height": int(height),
+        "payload": {
+            "files": [{"name": n, "b64": b} for n, b in zip(names, b64)],
+            "pose": base64.b64encode(pose.read_bytes()).decode() if pose else None,
+        },
+        "origin": f"{url.scheme}://{url.netloc}",
+        "src": f"{url.geturl().rstrip('/')}/index.html?embed=1",
+    }
+    return _SCRIPT % {
+        k: json.dumps(v).replace("<", "\\u003c") for k, v in literals.items()
+    }
+
+
+def _recording(dataset, index: int) -> Path:
+    """File behind ``dataset.datasets[index]``: the one its mne ``raw`` reads (a
+    data file maps back to its header; lazily downloading datasets fetch it when
+    ``raw`` is accessed) or the recorded ``description["path"]``."""
+    ds = dataset.datasets[index]
+    names = [
+        Path(f)
+        for f in getattr(getattr(ds, "raw", None), "filenames", None) or ()
+        if isinstance(f, (str, os.PathLike))
+    ]
+    if len(names) > 1:
+        raise ValueError(
+            f"{type(dataset).__name__}[{index}]: split recordings are not supported"
+        )
+    desc = getattr(ds, "description", None)  # dict or pandas Series
+    recorded = desc.get("path") if desc is not None else None
+    path = (
+        names[0]
+        if names
+        else Path(recorded)
+        if isinstance(recorded, (str, os.PathLike))
+        else None
+    )
+    if path is None:
+        raise ValueError(
+            f"{type(dataset).__name__}[{index}] is not backed by a recording file"
+        )
+    header = path.with_suffix(_HEADER.get(path.suffix.lower(), path.suffix))
+    return header if header.is_file() else path
+
+
+def plot(
+    dataset,
+    index: int = 0,
+    *,
+    height: int = 520,
+    cdn_url: str = CDN,
+    max_bytes: int = MAX_BYTES,
+):
+    """Show one recording in the eegdash-viewer inside a Jupyter cell.
+
+    Serverless: the recording bytes (as on disk) are inlined in the output and
+    pushed to the viewer at ``cdn_url`` over ``postMessage``; the output renders
+    when the cell ran in your session or the notebook is trusted. A
+    ``*_desc-pose.json`` sidecar next to the recording adds the synchronized
+    hand-pose panel. Needs IPython (soft dependency).
+
+    Parameters
+    ----------
+    index : int
+        Recording to display.
+    height : int
+        Viewer height in pixels.
+    cdn_url : str
+        Base URL of a deployed eegdash-viewer.
+    max_bytes : int
+        Refuse to inline more than this much base64 (default 64 MiB); the
+        payload is saved with the notebook and, like any cell output, stays
+        referenced by IPython's ``Out`` history for the session.
+
+    Returns
+    -------
+    IPython.display.HTML
+    """
+    ipython = _soft_import("IPython", purpose=f"{type(dataset).__name__}.plot()")
+    return ipython.display.HTML(
+        build_viewer_html(
+            _recording(dataset, index),
+            height=height,
+            cdn_url=cdn_url,
+            max_bytes=max_bytes,
+        )
+    )

{braindecode-1.8.0.dev1128 → braindecode-1.8.0.dev169056309}/braindecode/datasets/base.py

@@ -32,6 +32,7 @@ from mne.utils.docs import deprecated
 from torch.utils.data import ConcatDataset, Dataset, IterableDataset
 from typing_extensions import TypeVar
 
+from ._notebook_viewer import plot as _viewer_plot
 from .bids.hub import HubDatasetMixin
 from .bids.hub_io import _restore_nan_from_json
 from .registry import register_dataset
@@ -59,6 +60,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",
@@ -1173,6 +1175,8 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
         If True, defer computing cumulative sizes until length or item access.
     """
 
+    plot = _viewer_plot  # eegdash-viewer embed; a direct member so the API docs list it
+
     datasets: list[T]
 
     def __init__(
@@ -1374,6 +1378,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.dev1128 → braindecode-1.8.0.dev169056309}/braindecode/functional/functions.py

@@ -40,7 +40,7 @@ def drop_path(
     x : torch.Tensor
         input tensor
     drop_prob : float, optional
-        survival rate (i.e. probability of being kept), by default 0.0
+        probability of dropping a path, by default 0.0
     training : bool, optional
         whether the model is in training mode, by default False
     scale_by_keep : bool, optional
@@ -68,9 +68,10 @@ def drop_path(
     shape = (x.shape[0],) + (1,) * (
         x.ndim - 1
     )  # work with diff dim tensors, not just 2D ConvNets
-    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    probabilities = torch.full(shape, keep_prob, dtype=torch.float32, device=x.device)
+    random_tensor = torch.bernoulli(probabilities).to(dtype=x.dtype)
     if keep_prob > 0.0 and scale_by_keep:
-        random_tensor.div_(keep_prob)
+        random_tensor = random_tensor / keep_prob
     return x * random_tensor
 
 

{braindecode-1.8.0.dev1128 → braindecode-1.8.0.dev169056309}/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.dev1128 → braindecode-1.8.0.dev169056309}/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,