Braindecode 1.3.2.dev179983996__tar.gz → 1.3.2.dev181313884__tar.gz

{braindecode-1.3.2.dev179983996/braindecode.egg-info → braindecode-1.3.2.dev181313884}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev179983996
+Version: 1.3.2.dev181313884
 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.2
+Requires-Dist: torch>=2.0
 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: pandas<3.0.0
+Requires-Dist: pydantic>=2.0
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.4.3; extra == "moabb"
 Provides-Extra: eegprep

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev181313884}/braindecode/datasets/bids/hub.py

@@ -845,11 +845,16 @@ 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,
+                        metadata["i_start_in_trial"].values.astype(int),
                         np.zeros(len(metadata), dtype=int),
-                        metadata["target"].values,
+                        event_ids,
                     ]
                 )
                 epochs = mne.EpochsArray(data, info, events=events, metadata=metadata)

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev181313884}/braindecode/models/__init__.py

@@ -55,6 +55,7 @@ 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,
 )
@@ -118,6 +119,7 @@ __all__ = [
     "TIDNet",
     "TSception",
     "USleep",
+    "build_model_config",
     "_init_models_dict",
     "models_mandatory_parameters",
 ]

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev181313884}/braindecode/models/base.py

@@ -17,6 +17,13 @@ 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(
@@ -52,7 +59,25 @@ def deprecated_args(obj, *old_new_args):
     return out_args
 
 
-class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta):
+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):
     """
     Mixin class for all EEG models in braindecode.
 
@@ -86,59 +111,66 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
     -----
     If some input signal-related parameters are not specified,
     there will be an attempt to infer them from the other parameters.
+    """
 
+    #: Template for model-specific Hub integration notes appended by
+    #: ``__init_subclass__``.  ``{name}`` is replaced with the concrete
+    #: model class name.
+    _HUB_NOTES_TEMPLATE = """
     .. rubric:: Hugging Face Hub integration
 
     When the optional ``huggingface_hub`` package is installed, all models
     automatically gain the ability to be pushed to and loaded from the
     Hugging Face Hub. Install with::
 
-        pip install braindecode[hug]
+        pip install braindecode[hub]
 
     **Pushing a model to the Hub:**
 
-    .. code-block:: python
+    .. code-block::
 
-        from braindecode.models import EEGNetv4
+        from braindecode.models import {name}
 
         # Train your model
-        model = EEGNetv4(n_chans=22, n_outputs=4, n_times=1000)
+        model = {name}(n_chans=22, n_outputs=4, n_times=1000)
         # ... training code ...
 
         # Push to the Hub
         model.push_to_hub(
-            repo_id="username/my-eegnet-model", commit_message="Initial model upload"
+            repo_id="username/my-{name_lower}-model",
+            commit_message="Initial model upload",
         )
 
     **Loading a model from the Hub:**
 
-    .. code-block:: python
+    .. code-block::
 
-        from braindecode.models import EEGNetv4
+        from braindecode.models import {name}
 
         # Load pretrained model
-        model = EEGNetv4.from_pretrained("username/my-eegnet-model")
+        model = {name}.from_pretrained("username/my-{name_lower}-model")
 
     The integration automatically handles EEG-specific parameters (n_chans,
     n_times, sfreq, chs_info, etc.) by saving them in a config file alongside
     the model weights. This ensures that loaded models are correctly configured
     for their original data specifications.
 
-    .. 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.
-
-        To use non-default model parameters, specify them explicitly when calling
-        :func:`from_pretrained()`::
-
-            model = EEGNet.from_pretrained("user/model", dropout=0.3, activation='relu')
-
-        Full parameter serialization will be addressed in a future update.
+    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.
     """
 
     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.
+        if cls.__doc__ is not None:
+            hub_notes = cls._HUB_NOTES_TEMPLATE.format(
+                name=cls.__name__,
+                name_lower=cls.__name__.lower(),
+            )
+            cls.__doc__ = cls.__doc__.rstrip() + "\n" + hub_notes
+
         if not HAS_HF_HUB:
             super().__init_subclass__(**kwargs)
             return
@@ -157,6 +189,19 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
         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,
@@ -164,6 +209,7 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
             repo_url=repo_url,
             library_name=library_name,
             license=license,
+            coders=coders,
             **kwargs,
         )
 
@@ -205,6 +251,17 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
         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
@@ -315,6 +372,104 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
                     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)
+
+    def reset_head(self, n_outputs):
+        """Replace the classification head for a new number of outputs.
+
+        This is called automatically by :meth:`from_pretrained` when the
+        user passes an ``n_outputs`` that differs from the saved config.
+        Override in subclasses that need a model-specific head structure.
+
+        Parameters
+        ----------
+        n_outputs : int
+            New number of output classes.
+
+        Examples
+        --------
+        >>> from braindecode.models import BENDR
+        >>> model = BENDR(n_chans=22, n_times=1000, n_outputs=4)
+        >>> model.reset_head(10)
+        >>> model.n_outputs
+        10
+
+        .. versionadded:: 1.4
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not implement reset_head(). "
+            "Override this method to support changing n_outputs after "
+            "loading pretrained weights."
+        )
+
     mapping: Optional[Dict[str, str]] = None
 
     def load_state_dict(self, state_dict, *args, **kwargs):
@@ -419,54 +574,22 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
 
     @staticmethod
     def _serialize_chs_info(chs_info):
-        """
-        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.
-        """
+        """Serialize MNE channel info (``info["chs"]``) to JSON-compatible dicts."""
         if chs_info is None:
             return None
-
+        _INT_FIELDS = ("kind", "coil_type", "unit")
+        _FLOAT_FIELDS = ("cal", "range")
         serialized = []
         for ch in chs_info:
-            # 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
+            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)
             if "loc" in ch and ch["loc"] is not None:
                 ch_dict["loc"] = (
                     ch["loc"].tolist()
@@ -474,35 +597,19 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
                     else list(ch["loc"])
                 )
             serialized.append(ch_dict)
-
         return serialized
 
     @staticmethod
     def _deserialize_chs_info(chs_info_dict):
-        """
-        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.
-        """
+        """Deserialize JSON channel dicts back to MNE-compatible format."""
         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):
@@ -522,16 +629,8 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
 
         save_directory = Path(save_directory)
 
-        # 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__,
-        }
+        config = build_model_config(self)
+        config["braindecode_version"] = __version__
 
         # Save to config.json
         config_path = save_directory / "config.json"
@@ -570,7 +669,42 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
             **model_kwargs,
         ):
             model_kwargs.pop("braindecode_version", None)
-            return super()._from_pretrained(  # type: ignore
+            resolve_type_kwargs(cls, model_kwargs)
+
+            # Read saved n_outputs from config.json to detect when the
+            # user wants a different number of outputs.  Works for both
+            # local directories and Hub repo IDs.
+            saved_n_outputs = None
+            try:
+                if Path(model_id).is_dir():
+                    config_file = Path(model_id) / "config.json"
+                else:
+                    config_file = huggingface_hub.hf_hub_download(
+                        repo_id=model_id,
+                        filename="config.json",
+                        revision=revision,
+                        cache_dir=cache_dir,
+                        force_download=force_download,
+                        token=token,
+                        local_files_only=local_files_only,
+                    )
+                with open(config_file, "r") as f:
+                    saved_n_outputs = json.load(f).get("n_outputs")
+            except (OSError, json.JSONDecodeError, KeyError):
+                pass  # config unavailable; skip reset_head logic
+
+            requested_n_outputs = model_kwargs.get("n_outputs")
+
+            # If the user requests different n_outputs, load with the
+            # saved value first (so weights match), then swap the head.
+            if (
+                saved_n_outputs is not None
+                and requested_n_outputs is not None
+                and requested_n_outputs != saved_n_outputs
+            ):
+                model_kwargs["n_outputs"] = saved_n_outputs
+
+            model = super()._from_pretrained(  # type: ignore
                 model_id=model_id,
                 revision=revision,
                 cache_dir=cache_dir,
@@ -581,3 +715,20 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
                 strict=strict,
                 **model_kwargs,
             )
+
+            if (
+                saved_n_outputs is not None
+                and requested_n_outputs is not None
+                and requested_n_outputs != saved_n_outputs
+            ):
+                try:
+                    model.reset_head(requested_n_outputs)
+                except NotImplementedError:
+                    raise ValueError(
+                        f"{type(model).__name__} does not support changing "
+                        f"n_outputs after loading. Saved model has "
+                        f"n_outputs={saved_n_outputs}, but "
+                        f"n_outputs={requested_n_outputs} was requested."
+                    ) from None
+
+            return model

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev181313884}/braindecode/models/bendr.py

@@ -291,15 +291,23 @@ class BENDR(EEGModuleMixin, nn.Module):
         else:
             in_features = encoder_h
 
+        self._head_in_features = in_features
         self.final_layer = None
         if self.include_final_layer:
-            # in_features: encoder_h*4 (encoder_only) or encoder_h (full model)
-            linear = nn.Linear(in_features=in_features, out_features=self.n_outputs)
-            self.final_layer = nn.utils.parametrizations.weight_norm(
-                linear, name="weight", dim=1
-            )
+            self._build_head(self.n_outputs)
 
-    def forward(self, x):
+    def _build_head(self, n_outputs):
+        linear = nn.Linear(in_features=self._head_in_features, out_features=n_outputs)
+        self.final_layer = nn.utils.parametrizations.weight_norm(
+            linear, name="weight", dim=1
+        )
+
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        self.include_final_layer = True
+        self._build_head(n_outputs)
+
+    def forward(self, x, return_features=False):
         if self.channel_projection is not None:
             x = self.channel_projection(x)
         encoded = self.encoder(x)
@@ -328,6 +336,9 @@ class BENDR(EEGModuleMixin, nn.Module):
             feature = context[:, :, 0]
             # feature: [batch_size, encoder_h]
 
+        if return_features:
+            return {"features": feature, "cls_token": None}
+
         if self.final_layer is not None:
             feature = self.final_layer(feature)
             # feature: [batch_size, n_outputs]

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev181313884}/braindecode/models/biot.py

@@ -183,13 +183,22 @@ class BIOT(EEGModuleMixin, nn.Module):
             attn_layer_dropout=att_layer_drop_prob,
         )
 
+        self._head_activation = activation
         self.final_layer = _ClassificationHead(
             emb_size=self.embed_dim,
             n_outputs=self.n_outputs,
             activation=activation,
         )
 
-    def forward(self, x):
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        self.final_layer = _ClassificationHead(
+            emb_size=self.embed_dim,
+            n_outputs=n_outputs,
+            activation=self._head_activation,
+        )
+
+    def forward(self, x, return_features=False):
         """
         Pass the input through the BIOT encoder, and then through the
         classification head.
@@ -198,15 +207,24 @@ class BIOT(EEGModuleMixin, nn.Module):
         ----------
         x: Tensor
             (batch_size, n_channels, n_times)
+        return_features : bool
+            If True, return a dict with ``"features"`` and ``"cls_token"``
+            instead of the classification output.
 
         Returns
         -------
-        out: Tensor
-            (batch_size, n_outputs)
-        (out, emb): tuple Tensor
-            (batch_size, n_outputs), (batch_size, emb_size)
+        torch.Tensor or tuple or dict
+            Default: ``torch.Tensor`` of shape ``(batch_size, n_outputs)``.
+            If ``return_features=True``: ``dict`` with ``"features"``
+            ``(batch_size, emb_size)`` and ``"cls_token"`` (``None``).
+            If legacy ``return_feature=True`` (init param):
+            ``(out, emb)`` tuple (ignored when ``return_features=True``).
         """
         emb = self.encoder(x)
+
+        if return_features:
+            return {"features": emb, "cls_token": None}
+
         x = self.final_layer(emb)
 
         if self.return_feature:

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev181313884}/braindecode/models/cbramod.py

@@ -205,13 +205,33 @@ class CBraMod(EEGModuleMixin, nn.Module):
         self.encoder = _TransformerEncoder(encoder_layer, num_layers=n_layer)
         self.proj_out = nn.Sequential(nn.Linear(d_model, emb_dim))
 
+        self._emb_dim = emb_dim
+        self._patch_size = patch_size
         self._weights_init()
 
-        self.final_layer = (
-            nn.Identity()
-            if return_encoder_output
-            else nn.Sequential(nn.Flatten(), nn.LazyLinear(self.n_outputs))
-        )
+        if return_encoder_output:
+            self.final_layer = nn.Identity()
+        elif self._n_times is not None and self._n_chans is not None:
+            n_patch = self._n_times // patch_size
+            flat_dim = self._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)
+            )
+
+    def reset_head(self, n_outputs):
+        self._n_outputs = n_outputs
+        if self._n_times is not None and self._n_chans is not None:
+            n_patch = self._n_times // self._patch_size
+            flat_dim = self._n_chans * n_patch * self._emb_dim
+            self.final_layer = nn.Sequential(
+                nn.Flatten(), nn.Linear(flat_dim, n_outputs)
+            )
+        else:
+            self.final_layer = nn.Sequential(nn.Flatten(), nn.LazyLinear(n_outputs))
 
     def _weights_init(self):
         for m in self.modules():
@@ -223,11 +243,13 @@ class CBraMod(EEGModuleMixin, nn.Module):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 
-    def forward(self, x, mask=None):
+    def forward(self, x, mask=None, return_features=False):
         x = self.rearrange(x)
         patch_emb = self.patch_embedding(x, mask)
         feats = self.encoder(patch_emb)
         out = self.proj_out(feats)
+        if return_features:
+            return {"features": out, "cls_token": None}
         return self.final_layer(out)