Braindecode 1.3.2.dev181313884__tar.gz → 1.3.2.dev182544824__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev181313884
+Version: 1.3.2.dev182544824
 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.3.2.dev181313884 → braindecode-1.3.2.dev182544824}/braindecode/models/base.py

@@ -59,25 +59,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.
 
@@ -162,8 +144,8 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
 
     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__,
@@ -173,6 +155,7 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
 
         if not HAS_HF_HUB:
             super().__init_subclass__(**kwargs)
+            track_model_init_kwargs(cls)
             return
 
         base_tags = ["braindecode", cls.__name__]
@@ -212,6 +195,7 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
             coders=coders,
             **kwargs,
         )
+        track_model_init_kwargs(cls)
 
     def __init__(
         self,
@@ -442,34 +426,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
         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):
@@ -670,41 +626,7 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
         ):
             model_kwargs.pop("braindecode_version", None)
             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
+            return super()._from_pretrained(  # type: ignore
                 model_id=model_id,
                 revision=revision,
                 cache_dir=cache_dir,
@@ -715,20 +637,3 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
                 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.dev181313884 → braindecode-1.3.2.dev182544824}/braindecode/models/bendr.py

@@ -291,23 +291,15 @@ 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:
-            self._build_head(self.n_outputs)
-
-    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)
+            # 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
+            )
 
-    def forward(self, x, return_features=False):
+    def forward(self, x):
         if self.channel_projection is not None:
             x = self.channel_projection(x)
         encoded = self.encoder(x)
@@ -336,9 +328,6 @@ 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.dev181313884 → braindecode-1.3.2.dev182544824}/braindecode/models/biot.py

@@ -183,22 +183,13 @@ 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 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):
+    def forward(self, x):
         """
         Pass the input through the BIOT encoder, and then through the
         classification head.
@@ -207,24 +198,15 @@ 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
         -------
-        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``).
+        out: Tensor
+            (batch_size, n_outputs)
+        (out, emb): tuple Tensor
+            (batch_size, n_outputs), (batch_size, emb_size)
         """
         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.dev181313884 → braindecode-1.3.2.dev182544824}/braindecode/models/cbramod.py

@@ -205,15 +205,20 @@ 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()
 
+        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 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
+        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)
             )
@@ -222,17 +227,6 @@ class CBraMod(EEGModuleMixin, nn.Module):
                 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():
             if isinstance(m, nn.Linear):
@@ -243,13 +237,11 @@ class CBraMod(EEGModuleMixin, nn.Module):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 
-    def forward(self, x, mask=None, return_features=False):
+    def forward(self, x, mask=None):
         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)
 
 

{braindecode-1.3.2.dev181313884 → braindecode-1.3.2.dev182544824}/braindecode/models/eegpt.py

@@ -368,13 +368,6 @@ class EEGPT(EEGModuleMixin, nn.Module):
         """Number of temporal patches from encoder."""
         return self.target_encoder.num_patches[1]
 
-    def reset_head(self, n_outputs):
-        self._n_outputs = n_outputs
-        self.final_layer = _LinearConstraintProbe(
-            **self.get_probe_params(),
-            n_outputs=n_outputs,
-        )
-
     def get_probe_params(self) -> dict:
         """Get parameters needed to create a _LinearConstraintProbe.
 
@@ -389,7 +382,7 @@ class EEGPT(EEGModuleMixin, nn.Module):
             "embed_dim": self.embed_dim,
         }
 
-    def forward(self, x, return_features=False):
+    def forward(self, x):
         """
         Forward pass.
 
@@ -397,13 +390,10 @@ class EEGPT(EEGModuleMixin, nn.Module):
         ----------
         x : torch.Tensor
             EEG data of shape (batch, n_chans, n_times).
-        return_features : bool
-            If True, return a dict with ``"features"`` and ``"cls_token"``
-            instead of the classification output.
 
         Returns
         -------
-        torch.Tensor or dict
+        torch.Tensor
             Model output. Shape depends on `n_outputs` and `return_encoder_output`.
         """
         # Channel projection (if configured)
@@ -412,9 +402,6 @@ class EEGPT(EEGModuleMixin, nn.Module):
         # z shape: (batch, n_patches, embed_num, embed_dim)
         z = self.target_encoder(x, self.chans_id)
 
-        if return_features:
-            return {"features": z.flatten(2), "cls_token": None}
-
         if self.return_encoder_output:
             return z
 

{braindecode-1.3.2.dev181313884 → braindecode-1.3.2.dev182544824}/braindecode/models/labram.py

@@ -790,7 +790,6 @@ class Labram(EEGModuleMixin, nn.Module):
         ch_names: list[str] | None = None,
         return_patch_tokens=False,
         return_all_tokens=False,
-        return_features=False,
     ):
         """
         Forward the input EEG data through the model.
@@ -806,9 +805,6 @@ class Labram(EEGModuleMixin, nn.Module):
             If not provided, the channels provided during model initialization
             (via ``chs_info``), channels will be used instead
             If neither is provided, an error will be raised.
-        return_features : bool
-            If True, return a dict with ``"features"`` (patch tokens) and
-            ``"cls_token"`` instead of the classification output.
         return_patch_tokens: bool
             Return the patch tokens
         return_all_tokens: bool
@@ -816,18 +812,12 @@ class Labram(EEGModuleMixin, nn.Module):
 
         Returns
         -------
-        torch.Tensor or dict
+        torch.Tensor
             The output of the model with dimensions (batch, n_outputs)
         """
         # Apply automatic channel reordering if configured and input_chans not provided
         x, input_chans = self._select_channels(x, ch_names=ch_names)
 
-        if return_features:
-            x = self.forward_features(
-                x, input_chans=input_chans, return_all_tokens=True
-            )
-            return {"features": x[:, 1:, :], "cls_token": x[:, 0, :]}
-
         x = self.forward_features(
             x,
             input_chans=input_chans,
@@ -861,10 +851,6 @@ class Labram(EEGModuleMixin, nn.Module):
             nn.Linear(self.embed_dim, n_outputs) if n_outputs > 0 else nn.Identity()
         )
 
-    def reset_head(self, n_outputs):
-        self._n_outputs = n_outputs
-        self.reset_classifier(n_outputs)
-
     def _adj_temporal_embedding(self, num_ch, batch_size, dim_embed=None):
         """
         Adjust the dimensions of the time embedding to match the

{braindecode-1.3.2.dev181313884 → braindecode-1.3.2.dev182544824}/braindecode/models/reve.py

@@ -311,7 +311,19 @@ class REVE(EEGModuleMixin, nn.Module):
             geglu=self.use_geglu,
         )
 
-        self._build_head(self.n_outputs)
+        final_dim = self._get_flattened_output_dim()
+
+        if self.use_attention_pooling:
+            self.final_layer = nn.Sequential(
+                nn.LayerNorm(self.embed_dim),
+                nn.Linear(self.embed_dim, self.n_outputs),
+            )
+        else:
+            self.final_layer = nn.Sequential(
+                nn.Flatten(),
+                nn.LayerNorm(final_dim),
+                nn.Linear(final_dim, self.n_outputs),
+            )
 
         if self.use_attention_pooling:
             self.cls_query_token = nn.Parameter(torch.randn(1, 1, self.embed_dim))
@@ -340,24 +352,6 @@ class REVE(EEGModuleMixin, nn.Module):
         flat_dim = self.n_chans * n_patches * self.embed_dim
         return flat_dim
 
-    def _build_head(self, n_outputs):
-        if self.use_attention_pooling:
-            self.final_layer = nn.Sequential(
-                nn.LayerNorm(self.embed_dim),
-                nn.Linear(self.embed_dim, n_outputs),
-            )
-        else:
-            final_dim = self._get_flattened_output_dim()
-            self.final_layer = nn.Sequential(
-                nn.Flatten(),
-                nn.LayerNorm(final_dim),
-                nn.Linear(final_dim, n_outputs),
-            )
-
-    def reset_head(self, n_outputs):
-        self._n_outputs = n_outputs
-        self._build_head(n_outputs)
-
     def get_positions(self, channel_names: list[str]) -> torch.Tensor:
         """Fetch channel positions from the position bank. The position bank is downloaded when the model is instantiated.
 
@@ -379,8 +373,7 @@ class REVE(EEGModuleMixin, nn.Module):
         eeg: torch.Tensor,
         pos: Optional[torch.Tensor] = None,
         return_output: bool = False,
-        return_features: bool = False,
-    ) -> Union[torch.Tensor, list[torch.Tensor], dict]:
+    ) -> Union[torch.Tensor, list[torch.Tensor]]:
         """
         Forward pass of the model.
 
@@ -399,22 +392,16 @@ class REVE(EEGModuleMixin, nn.Module):
         return_output : bool, optional
             If True, returns the output from the transformer directly.
             If False, applies the final layer and returns the processed output. Default is False.
-        return_features : bool, optional
-            If True, returns a dict ``{"features": Tensor, "cls_token": None}``
-            with encoder features before the final layer. Default is False.
 
         Returns
         -------
-        torch.Tensor or list[torch.Tensor] or dict
-            Default: ``torch.Tensor`` after the final layer.
-            If ``return_output=True``: ``list[torch.Tensor]`` from transformer layers.
-            If ``return_features=True``: ``dict`` with ``"features"`` and ``"cls_token"`` keys.
-        """
+        Union[torch.Tensor, list[torch.Tensor]]
+            - If `return_output` is False: Returns a single `torch.Tensor` (output after final layer).
+            - If `return_output` is True: Returns a `list[torch.Tensor]` (outputs from transformer layers).
 
-        if return_output and return_features:
-            raise ValueError(
-                "return_output and return_features are mutually exclusive."
-            )
+            The output tensor(s) from the model. If `return_output` is True,
+            returns the transformer output; otherwise, returns the output after the final layer.
+        """
 
         patches = eeg.unfold(
             dimension=2,
@@ -466,9 +453,6 @@ class REVE(EEGModuleMixin, nn.Module):
             emb=self.embed_dim,
         )
 
-        if return_features:
-            return {"features": x, "cls_token": None}
-
         if self.use_attention_pooling:
             x = self._attention_pooling(x)
 

{braindecode-1.3.2.dev181313884 → braindecode-1.3.2.dev182544824}/braindecode/models/signal_jepa.py

@@ -222,13 +222,11 @@ class SignalJEPA(_BaseSignalJEPA):
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
         self.final_layer = nn.Identity()
 
-    def forward(self, X, ch_idxs: torch.Tensor | None = None, return_features=False):  # type: ignore
+    def forward(self, X, ch_idxs: torch.Tensor | None = None):  # type: ignore
         local_features = self.feature_encoder(X)  # type: ignore
         pos_encoding = self.pos_encoder(local_features, ch_idxs=ch_idxs)  # type: ignore
         local_features += pos_encoding  # type: ignore
         contextual_features = self.transformer.encoder(local_features)  # type: ignore
-        if return_features:
-            return {"features": contextual_features, "cls_token": None}
         y = self.final_layer(contextual_features)  # type: ignore
         return y  # type: ignore
 
@@ -315,8 +313,6 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
             _init_transformer=_init_transformer,
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
-        self._clf_conv_layers_spec = feature_encoder__conv_layers_spec
-        self._clf_n_spat_filters = n_spat_filters
         self.final_layer = _get_separable_clf_layer(
             conv_layers_spec=feature_encoder__conv_layers_spec,
             n_chans=self.n_chans,
@@ -325,26 +321,6 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
             n_spat_filters=n_spat_filters,
         )
 
-    def reset_head(self, n_outputs):
-        self._n_outputs = n_outputs
-        self.final_layer = _get_separable_clf_layer(
-            conv_layers_spec=self._clf_conv_layers_spec,
-            n_chans=self.n_chans,
-            n_times=self.n_times,
-            n_classes=n_outputs,
-            n_spat_filters=self._clf_n_spat_filters,
-        )
-
-    def forward(self, X, ch_idxs: torch.Tensor | None = None, return_features=False):  # type: ignore
-        local_features = self.feature_encoder(X)  # type: ignore
-        pos_encoding = self.pos_encoder(local_features, ch_idxs=ch_idxs)  # type: ignore
-        local_features += pos_encoding  # type: ignore
-        contextual_features = self.transformer.encoder(local_features)  # type: ignore
-        if return_features:
-            return {"features": contextual_features, "cls_token": None}
-        y = self.final_layer(contextual_features)  # type: ignore
-        return y  # type: ignore
-
     @classmethod
     def from_pretrained(
         cls,
@@ -419,6 +395,14 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
 
         return new_model
 
+    def forward(self, X, ch_idxs: torch.Tensor | None = None):  # type: ignore
+        local_features = self.feature_encoder(X)  # type: ignore
+        pos_encoding = self.pos_encoder(local_features, ch_idxs=ch_idxs)  # type: ignore
+        local_features += pos_encoding  # type: ignore
+        contextual_features = self.transformer.encoder(local_features)  # type: ignore
+        y = self.final_layer(contextual_features)  # type: ignore
+        return y  # type: ignore
+
 
 class SignalJEPA_PostLocal(_BaseSignalJEPA):
     r"""Post-local downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
@@ -501,8 +485,6 @@ class SignalJEPA_PostLocal(_BaseSignalJEPA):
             _init_transformer=False,
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
-        self._clf_conv_layers_spec = feature_encoder__conv_layers_spec
-        self._clf_n_spat_filters = n_spat_filters
         self.final_layer = _get_separable_clf_layer(
             conv_layers_spec=feature_encoder__conv_layers_spec,
             n_chans=self.n_chans,
@@ -511,16 +493,6 @@ class SignalJEPA_PostLocal(_BaseSignalJEPA):
             n_spat_filters=n_spat_filters,
         )
 
-    def reset_head(self, n_outputs):
-        self._n_outputs = n_outputs
-        self.final_layer = _get_separable_clf_layer(
-            conv_layers_spec=self._clf_conv_layers_spec,
-            n_chans=self.n_chans,
-            n_times=self.n_times,
-            n_classes=n_outputs,
-            n_spat_filters=self._clf_n_spat_filters,
-        )
-
     @classmethod
     def from_pretrained(
         cls,
@@ -577,10 +549,8 @@ class SignalJEPA_PostLocal(_BaseSignalJEPA):
         new_model.feature_encoder = deepcopy(feature_encoder)
         return new_model
 
-    def forward(self, X, return_features=False):
+    def forward(self, X):
         local_features = self.feature_encoder(X)
-        if return_features:
-            return {"features": local_features, "cls_token": None}
         y = self.final_layer(local_features)
         return y
 
@@ -701,21 +671,14 @@ class SignalJEPA_PreLocal(_BaseSignalJEPA):
             nn.Conv2d(1, n_spat_filters, (self.n_chans, 1)),
             Rearrange("b spat_filters 1 time -> b spat_filters time"),
         )
-        self._out_emb_dim = _get_out_emb_dim(
+        out_emb_dim = _get_out_emb_dim(
             conv_layers_spec=feature_encoder__conv_layers_spec,
             n_times=self.n_times,
             n_spat_filters=n_spat_filters,
         )
         self.final_layer = nn.Sequential(
             nn.Flatten(start_dim=1),
-            nn.Linear(self._out_emb_dim, self.n_outputs),
-        )
-
-    def reset_head(self, n_outputs):
-        self._n_outputs = n_outputs
-        self.final_layer = nn.Sequential(
-            nn.Flatten(start_dim=1),
-            nn.Linear(self._out_emb_dim, n_outputs),
+            nn.Linear(out_emb_dim, self.n_outputs),
         )
 
     @classmethod
@@ -774,11 +737,9 @@ class SignalJEPA_PreLocal(_BaseSignalJEPA):
         new_model.feature_encoder = deepcopy(feature_encoder)
         return new_model
 
-    def forward(self, X, return_features=False):
+    def forward(self, X):
         X = self.spatial_conv(X)
         local_features = self.feature_encoder(X)
-        if return_features:
-            return {"features": local_features, "cls_token": None}
         y = self.final_layer(local_features)
         return y
 

braindecode-1.3.2.dev182544824/braindecode/version.py

@@ -0,0 +1 @@
+__version__ = "1.3.2.dev182544824"

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev181313884
+Version: 1.3.2.dev182544824
 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.3.2.dev181313884 → braindecode-1.3.2.dev182544824}/docs/conf.py

@@ -349,6 +349,13 @@ html_theme_options = {
     "footer_start": ["copyright"],
     # 'pygment_light_style': 'default',
     "analytics": dict(google_analytics_id="G-7Q43R82K6D"),
+    "announcement": (
+        "<strong>Using Braindecode in academic work?</strong> "
+        "<a class='braindecode-announcement-cta' href='cite.html'>Cite Braindecode</a> "
+        "<span class='braindecode-announcement-secondary'>"
+        "DOI: <a href='https://doi.org/10.5281/zenodo.17699192'>10.5281/zenodo.17699192</a>"
+        "</span>"
+    ),
 }
 
 # The name of an image file (relative to this directory) to place at the top