Braindecode 1.3.2.dev183987088__tar.gz → 1.3.2.dev184103859__tar.gz

{braindecode-1.3.2.dev183987088/braindecode.egg-info → braindecode-1.3.2.dev184103859}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev183987088
+Version: 1.3.2.dev184103859
 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>
@@ -22,7 +22,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
 License-File: NOTICE.txt
-Requires-Dist: torch>=2.0
+Requires-Dist: torch>=2.2
 Requires-Dist: torchaudio>=2.0
 Requires-Dist: mne>=1.11.0
 Requires-Dist: pandas<3.0.0
@@ -35,7 +35,7 @@ Requires-Dist: wfdb
 Requires-Dist: linear_attention_transformer
 Requires-Dist: docstring_inheritance
 Requires-Dist: rotary_embedding_torch
-Requires-Dist: pydantic>=2.0
+Requires-Dist: pandas<3.0.0
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.4.3; extra == "moabb"
 Provides-Extra: eegprep

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/datasets/bids/hub.py

@@ -845,16 +845,11 @@ class HubDatasetMixin:
 
                 # Convert to MNE objects and create dataset
                 info = Info.from_json_dict(info_dict)
-                targets = metadata["target"].values
-                if np.issubdtype(targets.dtype, np.integer):
-                    event_ids = targets
-                else:
-                    event_ids = np.ones(len(metadata), dtype=int)
                 events = np.column_stack(
                     [
-                        metadata["i_start_in_trial"].values.astype(int),
+                        metadata["i_start_in_trial"].values,
                         np.zeros(len(metadata), dtype=int),
-                        event_ids,
+                        metadata["target"].values,
                     ]
                 )
                 epochs = mne.EpochsArray(data, info, events=events, metadata=metadata)

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/models/__init__.py

@@ -55,7 +55,6 @@ from .tsinception import TSception
 from .usleep import USleep
 from .util import (
     _init_models_dict,
-    build_model_config,
     extract_channel_locations_from_chs_info,
     models_mandatory_parameters,
 )
@@ -119,7 +118,6 @@ __all__ = [
     "TIDNet",
     "TSception",
     "USleep",
-    "build_model_config",
     "_init_models_dict",
     "models_mandatory_parameters",
 ]

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/models/base.py

@@ -17,13 +17,6 @@ from docstring_inheritance import NumpyDocstringInheritanceInitMeta
 from mne.utils import _soft_import
 from torchinfo import ModelStatistics, summary
 
-from braindecode.models.util import (
-    _EEG_PARAMS,
-    _IMPORT_ADAPTER,
-    build_model_config,
-    resolve_type_kwargs,
-    track_model_init_kwargs,
-)
 from braindecode.version import __version__
 
 huggingface_hub = _soft_import(
@@ -59,25 +52,7 @@ def deprecated_args(obj, *old_new_args):
     return out_args
 
 
-class _BraindecodeDocstringMeta(NumpyDocstringInheritanceInitMeta):
-    """Defer ``__init__`` wrapping until after docstring inheritance.
-
-    ``NumpyDocstringInheritanceInitMeta`` uses ``inspect.unwrap()``
-    internally, which bypasses ``@wraps`` wrappers.  By wrapping
-    ``__init__`` *after* docstring processing, the metaclass sees the
-    unwrapped function and correctly inherits ``cls.__doc__``.
-    """
-
-    def __init__(cls, class_name, class_bases, class_dict):
-        super().__init__(class_name, class_bases, class_dict)
-        # Only wrap subclass __init__s, not EEGModuleMixin itself.
-        # Wrapping the mixin would cause super().__init__() calls to
-        # overwrite _braindecode_init_kwargs captured by the subclass.
-        if any(isinstance(b, _BraindecodeDocstringMeta) for b in class_bases):
-            track_model_init_kwargs(cls)
-
-
-class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
+class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta):
     """
     Mixin class for all EEG models in braindecode.
 
@@ -155,15 +130,19 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
     the model weights. This ensures that loaded models are correctly configured
     for their original data specifications.
 
-    All model parameters (both EEG-specific and model-specific such as
-    dropout rates, activation functions, number of filters) are automatically
-    saved to the Hub and restored when loading.
+    .. important::
+        Currently, only EEG-specific parameters (n_outputs, n_chans, n_times,
+        input_window_seconds, sfreq, chs_info) are saved to the Hub.
+        Model-specific parameters (e.g., dropout rates, activation functions,
+        number of filters) are not preserved and will use their default values
+        when loading from the Hub.
+        Full parameter serialization will be addressed in a future update.
     """
 
     def __init_subclass__(cls, **kwargs):
         # Append model-specific Hub integration notes to the docstring.
-        # This runs before the metaclass __init__, so the Hub notes will
-        # be included in the docstring that the metaclass processes.
+        # This runs after the metaclass, so we concatenate rather than
+        # override any existing Notes section in the subclass.
         if cls.__doc__ is not None:
             hub_notes = cls._HUB_NOTES_TEMPLATE.format(
                 name=cls.__name__,
@@ -189,19 +168,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
         repo_url = kwargs.pop("repo_url", "https://braindecode.org")
         library_name = kwargs.pop("library_name", "braindecode")
         license = kwargs.pop("license", "bsd-3-clause")
-
-        # Register a coder so that type[nn.Module] parameters
-        # (e.g. activation=nn.ELU) are serialized as importable
-        # strings in config.json and decoded back on load.
-        coders = kwargs.pop("coders", None) or {}
-        coders.setdefault(
-            type,
-            (
-                lambda t: f"{t.__module__}.{t.__qualname__}",
-                lambda data: _IMPORT_ADAPTER.validate_python(data),
-            ),
-        )
-
         # TODO: model_card_template can be added in the future for custom model cards
         super().__init_subclass__(
             tags=tags,
@@ -209,7 +175,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
             repo_url=repo_url,
             library_name=library_name,
             license=license,
-            coders=coders,
             **kwargs,
         )
 
@@ -251,17 +216,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
         self._n_times = n_times  # type: ignore[assignment]
         self._sfreq = sfreq  # type: ignore[assignment]
 
-        # Back-fill instance attributes from _hub_mixin_config for any
-        # params the subclass didn't store on self.  Skip EEG params
-        # (stored as self._*) and descriptors (properties).
-        for key, val in getattr(self, "_hub_mixin_config", {}).items():
-            if key in _EEG_PARAMS:
-                continue
-            if hasattr(getattr(type(self), key, None), "__get__"):
-                continue
-            if not hasattr(self, key):
-                setattr(self, key, val)
-
         super().__init__()
 
     @property
@@ -372,76 +326,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
                     raise ValueError(msg) from exc
                 raise exc
 
-    def get_config(self) -> dict:
-        """Return a JSON-serializable dict of all ``__init__`` parameters.
-
-        The returned dictionary can be saved to a JSON file and later
-        used with :meth:`from_config` to reconstruct the model (without
-        weights).  It is also used internally by :meth:`push_to_hub` to
-        persist the full model configuration.
-
-        Returns
-        -------
-        dict
-            All ``__init__`` parameters, JSON-serializable.
-            ``type[nn.Module]`` parameters (e.g. ``activation``) are
-            encoded as importable dotted-path strings.
-
-        Examples
-        --------
-        >>> import json
-        >>> from braindecode.models import EEGNet
-        >>> model = EEGNet(n_chans=22, n_times=1000, n_outputs=4, F1=16)
-        >>> config = model.get_config()
-        >>> config["F1"]
-        16
-        >>> # Save to disk
-        >>> with open("config.json", "w") as f:
-        ...     json.dump(config, f)
-
-        .. versionadded:: 1.4
-        """
-        return build_model_config(self)
-
-    @classmethod
-    def from_config(cls, config: dict) -> "EEGModuleMixin":
-        """Create a model instance from a configuration dict.
-
-        This is the inverse of :meth:`get_config`.  Weights are **not**
-        loaded -- use :meth:`from_pretrained` for that.
-
-        Parameters
-        ----------
-        config : dict
-            Configuration dict as returned by :meth:`get_config`.
-
-        Returns
-        -------
-        EEGModuleMixin
-            A new model instance.
-
-        Examples
-        --------
-        >>> import json
-        >>> from braindecode.models import EEGNet
-        >>> model = EEGNet(n_chans=22, n_times=1000, n_outputs=4, F1=16)
-        >>> config = model.get_config()
-        >>> # Reconstruct (without weights)
-        >>> model2 = EEGNet.from_config(config)
-        >>> model2.F1
-        16
-        >>> # Or from a JSON file
-        >>> with open("config.json") as f:
-        ...     config = json.load(f)
-        >>> model3 = EEGNet.from_config(config)
-
-        .. versionadded:: 1.4
-        """
-        config = dict(config)  # shallow copy
-        config.pop("braindecode_version", None)
-        resolve_type_kwargs(cls, config)
-        return cls(**config)
-
     mapping: Optional[Dict[str, str]] = None
 
     def load_state_dict(self, state_dict, *args, **kwargs):
@@ -546,22 +430,54 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
 
     @staticmethod
     def _serialize_chs_info(chs_info):
-        """Serialize MNE channel info (``info["chs"]``) to JSON-compatible dicts."""
+        """
+        Serialize MNE channel info to JSON-compatible format.
+
+        Parameters
+        ----------
+        chs_info : list of dict or None
+            Channel information from MNE Info object.
+
+        Returns
+        -------
+        list of dict or None
+            Serialized channel information that can be saved to JSON.
+        """
         if chs_info is None:
             return None
-        _INT_FIELDS = ("kind", "coil_type", "unit")
-        _FLOAT_FIELDS = ("cal", "range")
+
         serialized = []
         for ch in chs_info:
-            ch_dict = {"ch_name": ch.get("ch_name", "")}
-            for key in _INT_FIELDS:
-                val = ch.get(key)
-                if val is not None:
-                    ch_dict[key] = val if isinstance(val, str) else int(val)
-            for key in _FLOAT_FIELDS:
-                val = ch.get(key)
-                if val is not None:
-                    ch_dict[key] = float(val)
+            # Extract serializable fields from MNE channel info
+            ch_dict = {
+                "ch_name": ch.get("ch_name", ""),
+            }
+
+            # Handle kind field - can be either string or integer
+            kind_val = ch.get("kind")
+            if kind_val is not None:
+                ch_dict["kind"] = (
+                    kind_val if isinstance(kind_val, str) else int(kind_val)
+                )
+
+            # Add numeric fields with safe conversion
+            coil_type = ch.get("coil_type")
+            if coil_type is not None:
+                ch_dict["coil_type"] = int(coil_type)
+
+            unit = ch.get("unit")
+            if unit is not None:
+                ch_dict["unit"] = int(unit)
+
+            cal = ch.get("cal")
+            if cal is not None:
+                ch_dict["cal"] = float(cal)
+
+            range_val = ch.get("range")
+            if range_val is not None:
+                ch_dict["range"] = float(range_val)
+
+            # Serialize location array if present
             if "loc" in ch and ch["loc"] is not None:
                 ch_dict["loc"] = (
                     ch["loc"].tolist()
@@ -569,19 +485,35 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
                     else list(ch["loc"])
                 )
             serialized.append(ch_dict)
+
         return serialized
 
     @staticmethod
     def _deserialize_chs_info(chs_info_dict):
-        """Deserialize JSON channel dicts back to MNE-compatible format."""
+        """
+        Deserialize channel info from JSON-compatible format to MNE-like structure.
+
+        Parameters
+        ----------
+        chs_info_dict : list of dict or None
+            Serialized channel information.
+
+        Returns
+        -------
+        list of dict or None
+            Deserialized channel information compatible with MNE.
+        """
         if chs_info_dict is None:
             return None
+
         deserialized = []
         for ch_dict in chs_info_dict:
             ch = ch_dict.copy()
+            # Convert location back to numpy array if present
             if "loc" in ch and ch["loc"] is not None:
                 ch["loc"] = np.array(ch["loc"])
             deserialized.append(ch)
+
         return deserialized
 
     def _save_pretrained(self, save_directory):
@@ -601,8 +533,16 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
 
         save_directory = Path(save_directory)
 
-        config = build_model_config(self)
-        config["braindecode_version"] = __version__
+        # Collect EEG-specific configuration
+        config = {
+            "n_outputs": self._n_outputs,
+            "n_chans": self._n_chans,
+            "n_times": self._n_times,
+            "input_window_seconds": self._input_window_seconds,
+            "sfreq": self._sfreq,
+            "chs_info": self._serialize_chs_info(self._chs_info),
+            "braindecode_version": __version__,
+        }
 
         # Save to config.json
         config_path = save_directory / "config.json"
@@ -641,7 +581,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
             **model_kwargs,
         ):
             model_kwargs.pop("braindecode_version", None)
-            resolve_type_kwargs(cls, model_kwargs)
             return super()._from_pretrained(  # type: ignore
                 model_id=model_id,
                 revision=revision,

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/models/cbramod.py

@@ -207,25 +207,11 @@ class CBraMod(EEGModuleMixin, nn.Module):
 
         self._weights_init()
 
-        try:
-            n_times = self.n_times
-            n_chans = self.n_chans
-        except ValueError:
-            n_times = None
-            n_chans = None
-
-        if return_encoder_output:
-            self.final_layer = nn.Identity()
-        elif n_times is not None and n_chans is not None:
-            n_patch = n_times // patch_size
-            flat_dim = n_chans * n_patch * emb_dim
-            self.final_layer = nn.Sequential(
-                nn.Flatten(), nn.Linear(flat_dim, self.n_outputs)
-            )
-        else:
-            self.final_layer = nn.Sequential(
-                nn.Flatten(), nn.LazyLinear(self.n_outputs)
-            )
+        self.final_layer = (
+            nn.Identity()
+            if return_encoder_output
+            else nn.Sequential(nn.Flatten(), nn.LazyLinear(self.n_outputs))
+        )
 
     def _weights_init(self):
         for m in self.modules():

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/models/config.py

@@ -4,16 +4,13 @@ from types import UnionType
 from typing import Annotated, Any, Literal, Union, get_args, get_origin
 
 import numpy as np
-import pydantic
+from mne.utils import _soft_import
 from typing_extensions import TypedDict
 
 from braindecode.models.base import EEGModuleMixin
-from braindecode.models.util import (
-    SigArgName,
-    _init_models_dict,
-    models_dict,
-    models_mandatory_parameters,
-)
+from braindecode.models.util import SigArgName, models_dict, models_mandatory_parameters
+
+pydantic = _soft_import(name="pydantic", purpose="model configuration", strict=False)
 
 try:
     from numpydantic import NDArray, Shape
@@ -60,12 +57,8 @@ SIGNAL_ARGS_TYPES = {
 }
 
 
-class BaseBraindecodeModelConfig(pydantic.BaseModel):
-    """Base class for braindecode model pydantic configs."""
-
+class BaseBraindecodeModelConfig(pydantic.BaseModel):  # type: ignore
     def create_instance(self) -> EEGModuleMixin:
-        if self.model_name_ not in models_dict:
-            _init_models_dict()
         model_cls = models_dict[self.model_name_]
         kwargs = self.model_dump(mode="python", exclude={"model_name_"})
         if kwargs.get("n_chans") is not None and kwargs.get("chs_info") is not None:
@@ -97,6 +90,13 @@ def make_model_config(
     type
         A pydantic BaseModel subclass representing the model config.
     """
+    if not pydantic:
+        raise ImportError(
+            "pydantic is required to use make_model_config. "
+            "Please install braindecode[typing]."
+        )
+
+    # ironically, we need to ignore the type here to have the soft dependency.
 
     @pydantic.model_validator(mode="before")
     def validate_signal_params(cls, data: Any):
@@ -207,23 +207,23 @@ def make_model_config(
 # and define __all__ based on generated classes
 __all__ = ["make_model_config"]
 
-if not models_dict:
-    _init_models_dict()
-
-models_configs: list[type[BaseBraindecodeModelConfig]] = []
-for model_name, req, _ in models_mandatory_parameters:
-    model_cls = models_dict[model_name]
-    model_cfg = make_model_config(model_cls, req)
-    globals()[model_cfg.__name__] = model_cfg
-    __all__.append(model_cfg.__name__)
-    models_configs.append(model_cfg)
-
-BraindecodeModelConfig = Annotated[  # type: ignore[valid-type]
-    Union[tuple(models_configs)],
-    pydantic.Field(
-        discriminator="model_name_", description="Braindecode model configuration"
-    ),
-]
+if not pydantic:
+    pass
+else:
+    models_configs: list[type[BaseBraindecodeModelConfig]] = []
+    for model_name, req, _ in models_mandatory_parameters:
+        model_cls = models_dict[model_name]
+        model_cfg = make_model_config(model_cls, req)
+        globals()[model_cfg.__name__] = model_cfg
+        __all__.append(model_cfg.__name__)
+        models_configs.append(model_cfg)
+
+    BraindecodeModelConfig = Annotated[  # type: ignore
+        Union[tuple(models_configs)],
+        pydantic.Field(
+            discriminator="model_name_", description="Braindecode model configuration"
+        ),
+    ]
 
 # # Example usage:
 #

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/models/luna.py

@@ -7,7 +7,6 @@ Retrieved from https://openreview.net/forum?id=uazfjnFL0G
 
 Original Authors: Berkay Döner, Thorir Mar Ingolfsson
 Braindecode Adaptation: Bruno Aristimunha
-Contributions: Sarthak Tayal <sarthaktayal2@gmail.com>
 
 the LICENSE Of this file is APACHE-2.0.
 """
@@ -83,8 +82,11 @@ class LUNA(EEGModuleMixin, nn.Module):
 
            # Load pre-trained encoder weights
            state_dict = load_file(model_path)
-           # load_state_dict applies model.mapping automatically
-           model.load_state_dict(state_dict, strict=False)
+           # Apply key mapping for pretrained weights
+           mapping = model.mapping.copy()
+           mapping["cross_attn.temparature"] = "cross_attn.temperature"
+           mapped_state_dict = {mapping.get(k, k): v for k, v in state_dict.items()}
+           model.load_state_dict(mapped_state_dict, strict=False)
 
        To push your own trained model to the Hub:
 
@@ -168,7 +170,6 @@ class LUNA(EEGModuleMixin, nn.Module):
             "channel_location_embedder.0.fc2.bias": "channel_location_embedder.fc2.bias",
             "channel_location_embedder.0.norm.weight": "channel_location_embedder.norm.weight",
             "channel_location_embedder.0.norm.bias": "channel_location_embedder.norm.bias",
-            "cross_attn.temparature": "cross_attn.temperature",  # typo in pretrained weights
         }
 
         # Model parameters
@@ -314,10 +315,8 @@ class LUNA(EEGModuleMixin, nn.Module):
         channel_locations_emb = self.channel_location_embedder(channel_locations)
 
         x_tokenized = rearrange(x_masked, "B (C t) D -> (B t) C D", C=num_channels)
-        channel_locations_emb = (
-            channel_locations_emb.repeat_interleave(  # was repeat(), wrong dim ordering
-                num_patches_per_channel, dim=0
-            )
+        channel_locations_emb = channel_locations_emb.repeat(
+            num_patches_per_channel, 1, 1
         )
         x_tokenized = x_tokenized + channel_locations_emb
 

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/models/patchedtransformer.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import math
-import warnings
 from typing import Optional, Tuple
 
 import torch
@@ -497,10 +496,9 @@ class _TransformerEncoderLayer(nn.Module):
 
         if dim_feedforward is None:
             dim_feedforward = 4 * embed_dim
-            warnings.warn(
-                "dim_feedforward is set to 4*embed_dim, the default in "
-                "Vaswani et al. (Attention is all you need)",
-                stacklevel=2,
+            # note: preserve the original behaviour (print) from the provided code
+            print(
+                "dim_feedforward is set to 4*embed_dim, the default in Vaswani et al. (Attention is all you need)"
             )
 
         self.layer_norm_att = _LayerNorm(embed_dim, bias=bias)

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/models/reve.py

@@ -16,16 +16,7 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange
 from torch import nn
-
-# Safe import for older PyTorch versions (Support for Intel-based Macs)
-try:
-    from torch.nn.attention import SDPBackend, sdpa_kernel
-
-    HAS_SDPA = True
-except ImportError:
-    HAS_SDPA = False
-    SDPBackend = None
-    sdpa_kernel = None
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 from braindecode.models.base import EEGModuleMixin
 
@@ -539,22 +530,11 @@ class FeedForward(nn.Module):
 
 
 class ClassicalAttention(nn.Module):
-    def __init__(self, heads: int, use_sdpa: bool | None = None):
+    def __init__(self, heads: int, use_sdpa: bool = True):
         super().__init__()
+        self.use_sdpa = use_sdpa
         self.heads = heads
 
-        if use_sdpa is None:
-            self.use_sdpa = HAS_SDPA
-        elif use_sdpa is True and not HAS_SDPA:
-            logger.warning(
-                "SDPA (Scaled Dot Product Attention) was requested, but it is not "
-                "available in your current PyTorch version. Falling back to naive "
-                "implementation. Please upgrade to PyTorch >= 2.2 for SDPA support."
-            )
-            self.use_sdpa = False
-        else:
-            self.use_sdpa = use_sdpa
-
     def forward(self, qkv: torch.Tensor) -> torch.Tensor:
         # Split concatenated QKV into separate tensors
         # qkv shape: (batch, seq_len, 3 * heads * head_dim)

{braindecode-1.3.2.dev183987088 → braindecode-1.3.2.dev184103859}/braindecode/models/shallow_fbcsp.py

@@ -2,21 +2,19 @@
 #
 # License: BSD (3-clause)
 
-from __future__ import annotations
-
-from collections.abc import Callable
+from typing import Callable
 
 from einops.layers.torch import Rearrange
 from torch import nn
 from torch.nn import init
 
+from braindecode.functional import square
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import (
     CombinedConv,
     Ensure4d,
     Expression,
     SafeLog,
-    Square,
     SqueezeFinalOutput,
 )
 
@@ -47,14 +45,12 @@ class ShallowFBCSPNet(EEGModuleMixin, nn.Sequential):
     final_conv_length: int | str
         Length of the final convolution layer.
         If set to "auto", length of the input signal must be specified.
-    conv_nonlin: type[nn.Module] | Callable
-        Non-linear module class to be used after convolution layers.
-        For backward compatibility, callables are also accepted and wrapped
-        with :class:`~braindecode.modules.Expression`.
+    conv_nonlin: callable
+        Non-linear function to be used after convolution layers.
     pool_mode: str
         Method to use on pooling layers. "max" or "mean".
-    activation_pool_nonlin: type[nn.Module]
-        Non-linear module class to be used after pooling layers.
+    activation_pool_nonlin: callable
+        Non-linear function to be used after pooling layers.
     split_first_layer: bool
         Split first layer into temporal and spatial layers (True) or just use temporal (False).
         There would be no non-linearity between the split layers.
@@ -87,7 +83,7 @@ class ShallowFBCSPNet(EEGModuleMixin, nn.Sequential):
         pool_time_length=75,
         pool_time_stride=15,
         final_conv_length="auto",
-        conv_nonlin: type[nn.Module] | Callable = Square,
+        conv_nonlin: Callable = square,
         pool_mode="mean",
         activation_pool_nonlin: type[nn.Module] = SafeLog,
         split_first_layer=True,
@@ -167,11 +163,7 @@ class ShallowFBCSPNet(EEGModuleMixin, nn.Sequential):
                     n_filters_conv, momentum=self.batch_norm_alpha, affine=True
                 ),
             )
-        if isinstance(self.conv_nonlin, type):
-            self.add_module("conv_nonlin_exp", self.conv_nonlin())
-        else:
-            # Backward compat: accept a callable (e.g. square function)
-            self.add_module("conv_nonlin_exp", Expression(self.conv_nonlin))
+        self.add_module("conv_nonlin_exp", Expression(self.conv_nonlin))
         self.add_module(
             "pool",
             pool_class(