Repository: braindecode/braindecode

Issue:
CTNet: `RuntimeError: The size of tensor a (40) must match the size of tensor b (256) at non-singleton dimension 2`
The following works:
```python
from braindecode.models import CTNet

model = CTNet(
    n_chans=64,  # Number of EEG channels
    n_outputs=2,  # Number of output classes
    n_times=500,  # Length of the input sequence (e.g., 500 time steps)
    emb_size=40, # Default parameter value
)

model.to("cuda")
model.eval()

# Make a forward pass with dummy data
dummy_input = torch.randn(1, 64, 500).to("cuda")  # Batch size 1, 64 channels, 500 time steps
output = model(dummy_input)
```
But if `emb_size` is changed to, for example, `256`, it gives the following error:
```python
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[15], line 7
      4 # Make a forward pass with dummy data
      6 dummy_input = torch.randn(1, 64, 500).to("cuda")  # Batch size 1, 64 channels, 500 time steps
----> 7 output = model(dummy_input)
      8 print("Output shape:", output.shape)  # Should be [1, 2] for 2 classes
      9 print("Output[0, 0]:", output[0, 0])

File /envs/eeg/lib/python3.12/site-packages/torch/nn/modules/module.py:1739, in Module._wrapped_call_impl(self, *args, **kwargs)
   1737     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1738 else:
-> 1739     return self._call_impl(*args, **kwargs)

File /envs/eeg/lib/python3.12/site-packages/torch/nn/modules/module.py:1750, in Module._call_impl(self, *args, **kwargs)
   1745 # If we don't have any hooks, we want to skip the rest of the logic in
   1746 # this function, and just call forward.
   1747 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1748         or _global_backward_pre_hooks or _global_backward_hooks
   1749         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750     return forward_call(*args, **kwargs)
   1752 result = None
   1753 called_always_called_hooks = set()

File /envs/eeg/lib/python3.12/site-packages/braindecode/models/ctnet.py:206, in CTNet.forward(self, x)
    204 cnn = self.cnn(x)
    205 cnn = cnn * math.sqrt(self.emb_size)
--> 206 cnn = self.position(cnn)
    207 trans = self.trans(cnn)
    208 features = cnn + trans

File /envs/eeg/lib/python3.12/site-packages/torch/nn/modules/module.py:1739, in Module._wrapped_call_impl(self, *args, **kwargs)
   1737     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1738 else:
-> 1739     return self._call_impl(*args, **kwargs)

File /envs/eeg/lib/python3.12/site-packages/torch/nn/modules/module.py:1750, in Module._call_impl(self, *args, **kwargs)
   1745 # If we don't have any hooks, we want to skip the rest of the logic in
   1746 # this function, and just call forward.
   1747 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1748         or _global_backward_pre_hooks or _global_backward_hooks
   1749         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750     return forward_call(*args, **kwargs)
   1752 result = None
   1753 called_always_called_hooks = set()

File /envs/eeg/lib/python3.12/site-packages/braindecode/models/ctnet.py:449, in _PositionalEncoding.forward(self, x)
    447 seq_length = x.size(1)
    448 encoding = self.encoding[:, :seq_length, :]
--> 449 x = x + encoding
    450 return self.dropout(x)

RuntimeError: The size of tensor a (40) must match the size of tensor b (256) at non-singleton dimension 2
```


Discussion:
Okay... I see what you are doing... You are basically testing the braindecode's models to extract different embeddings and facing different issues because the braindecode models are not designed to be an embedding model's... I will thinking some better test and solve this issue across models
I am trying them out both as standalone models for classification, as well as feature extractors (removing the final layer) for a regression task :)
If you initialise everything with EEGRegressor and EEGClassifier, you will have the nn.Module that you need inside the variable ._module variables. The only thing is that you need to pass everything related to the dataset, signal-related parameters.

issue related....

https://github.com/braindecode/braindecode/issues/684

But the models need to work by themselves without the skorch wrapper, so the issue is welcome. Do you have more models that you are focusing?


I could not reproduce your issue with braindecode source:

https://colab.research.google.com/drive/18xGk90UtMi07f-LKSfpMf9bEfN4JIuQM?usp=sharing
What do you mean? `emb_size=40` is the default, changing that to another value raises the error:

<img width="1061" height="620" alt="Image" src="https://github.com/user-attachments/assets/fff2cc01-c925-41a8-be53-078ad3dd8ea4" />
Sorry, I misunderstood the issue. I am solving rn



Patch:
diff --git a/braindecode/models/ctnet.py b/braindecode/models/ctnet.py
index 7850bea9e..11acf3708 100644
--- a/braindecode/models/ctnet.py
+++ b/braindecode/models/ctnet.py
@@ -10,6 +10,7 @@
 from __future__ import annotations
 
 import math
+from typing import Optional
 
 import torch
 from einops.layers.torch import Rearrange
@@ -57,7 +58,7 @@ class CTNet(EEGModuleMixin, nn.Module):
         Activation function to use in the network.
     heads : int, default=4
         Number of attention heads in the Transformer encoder.
-    emb_size : int, default=40
+    emb_size : int or None, default=None
         Embedding size (dimensionality) for the Transformer encoder.
     depth : int, default=6
         Number of encoder layers in the Transformer.
@@ -110,11 +111,11 @@ def __init__(
         drop_prob_final: float = 0.5,
         # other parameters
         heads: int = 4,
-        emb_size: int = 40,
+        emb_size: Optional[int] = 40,
         depth: int = 6,
-        n_filters_time: int = 20,
+        n_filters_time: Optional[int] = None,
         kernel_size: int = 64,
-        depth_multiplier: int = 2,
+        depth_multiplier: Optional[int] = 2,
         pool_size_1: int = 8,
         pool_size_2: int = 8,
     ):
@@ -128,21 +129,18 @@ def __init__(
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
 
-        self.emb_size = emb_size
         self.activation_patch = activation_patch
         self.activation_transformer = activation_transformer
-
-        self.n_filters_time = n_filters_time
         self.drop_prob_cnn = drop_prob_cnn
         self.pool_size_1 = pool_size_1
         self.pool_size_2 = pool_size_2
-        self.depth_multiplier = depth_multiplier
         self.kernel_size = kernel_size
         self.drop_prob_posi = drop_prob_posi
         self.drop_prob_final = drop_prob_final
-
+        self.heads = heads
+        self.depth = depth
         # n_times - pool_size_1 / p
-        sequence_length = math.floor(
+        self.sequence_length = math.floor(
             (
                 math.floor((self.n_times - self.pool_size_1) / self.pool_size_1 + 1)
                 - self.pool_size_2
@@ -151,6 +149,10 @@ def __init__(
             + 1
         )
 
+        self.depth_multiplier, self.n_filters_time, self.emb_size = self._resolve_dims(
+            depth_multiplier, n_filters_time, emb_size
+        )
+
         # Layers
         self.ensuredim = Rearrange("batch nchans time -> batch 1 nchans time")
         self.flatten = nn.Flatten()
@@ -167,14 +169,17 @@ def __init__(
         )
 
         self.position = _PositionalEncoding(
-            emb_size=emb_size,
+            emb_size=self.emb_size,
             drop_prob=self.drop_prob_posi,
             n_times=self.n_times,
             pool_size=self.pool_size_1,
         )
 
         self.trans = _TransformerEncoder(
-            heads, depth, emb_size, activation=self.activation_transformer
+            self.heads,
+            self.depth,
+            self.emb_size,
+            activation=self.activation_transformer,
         )
 
         self.flatten_drop_layer = nn.Sequential(
@@ -183,7 +188,8 @@ def __init__(
         )
 
         self.final_layer = nn.Linear(
-            in_features=emb_size * sequence_length, out_features=self.n_outputs
+            in_features=int(self.emb_size * self.sequence_length),
+            out_features=self.n_outputs,
         )
 
     def forward(self, x: Tensor) -> Tensor:
@@ -210,6 +216,86 @@ def forward(self, x: Tensor) -> Tensor:
         out = self.final_layer(flatten_feature)
         return out
 
+    @staticmethod
+    def _resolve_dims(
+        depth_multiplier: Optional[int],
+        n_filters_time: Optional[int],
+        emb_size: Optional[int],
+    ) -> tuple[int, int, int]:
+        # Basic type/positivity checks for provided values
+        for name, val in (
+            ("depth_multiplier", depth_multiplier),
+            ("n_filters_time", n_filters_time),
+            ("emb_size", emb_size),
+        ):
+            if val is not None:
+                if not isinstance(val, int):
+                    raise TypeError(f"{name} must be int, got {type(val).__name__}")
+                if val <= 0:
+                    raise ValueError(f"{name} must be > 0, got {val}")
+
+        missing = [
+            k
+            for k, v in {
+                "depth_multiplier": depth_multiplier,
+                "n_filters_time": n_filters_time,
+                "emb_size": emb_size,
+            }.items()
+            if v is None
+        ]
+
+        if len(missing) >= 2:
+            # Too many unknowns → ambiguous
+            raise ValueError(
+                "Specify exactly two of {depth_multiplier, n_filters_time, emb_size}; the third will be inferred."
+            )
+
+        if len(missing) == 1:
+            # Infer the missing one
+            if missing[0] == "emb_size":
+                assert depth_multiplier is not None and n_filters_time is not None
+                emb_size = depth_multiplier * n_filters_time
+            elif missing[0] == "n_filters_time":
+                assert emb_size is not None and depth_multiplier is not None
+                if emb_size % depth_multiplier != 0:
+                    raise ValueError(
+                        f"emb_size={emb_size} must be divisible by depth_multiplier={depth_multiplier}"
+                    )
+                n_filters_time = emb_size // depth_multiplier
+            else:  # missing depth_multiplier
+                assert emb_size is not None and n_filters_time is not None
+                if emb_size % n_filters_time != 0:
+                    raise ValueError(
+                        f"emb_size={emb_size} must be divisible by n_filters_time={n_filters_time}"
+                    )
+                depth_multiplier = emb_size // n_filters_time
+
+        else:
+            # All provided: enforce consistency
+            assert (
+                depth_multiplier is not None
+                and n_filters_time is not None
+                and emb_size is not None
+            )
+            prod = depth_multiplier * n_filters_time
+            if prod != emb_size:
+                raise ValueError(
+                    "`depth_multiplier * n_filters_time` must equal `emb_size`, "
+                    f"but got {depth_multiplier} * {n_filters_time} = {prod} != {emb_size}. "
+                    "Fix by setting one of: "
+                    f"emb_size={prod}, "
+                    f"n_filters_time={emb_size // depth_multiplier if emb_size % depth_multiplier == 0 else 'not integer'}, "
+                    f"depth_multiplier={emb_size // n_filters_time if emb_size % n_filters_time == 0 else 'not integer'}."
+                )
+
+        # Ensure plain ints for the return type
+        assert (
+            depth_multiplier is not None
+            and n_filters_time is not None
+            and emb_size is not None
+        )
+        return depth_multiplier, n_filters_time, emb_size
+
 
 class _PatchEmbeddingEEGNet(nn.Module):
     def __init__(
diff --git a/docs/whats_new.rst b/docs/whats_new.rst
index 00cbd8ff8..13db51cae 100644
--- a/docs/whats_new.rst
+++ b/docs/whats_new.rst
@@ -29,6 +29,7 @@ Bugs
 - Fix initialization of :class:`braindecode.models.Deep4Net` (:gh:`752` by `Bruno Aristimunha`_)
 - Fix initialization of :class:`braindecode.models.ATCNet` (:gh:`754` by `Bruno Aristimunha`_)
 - Fix Labram model for small input (:gh:`757` by `Bruno Aristimunha`_)
+- Fix the CTNet parameters to make the embedding configuration easy (:gh:`760` by `Bruno Aristimunha`_)
 
 API changes
 ~~~~~~~~~~~


Test patch:
diff --git a/test/unit_tests/models/test_integration.py b/test/unit_tests/models/test_integration.py
index 82a9a87be..854ef73c4 100644
--- a/test/unit_tests/models/test_integration.py
+++ b/test/unit_tests/models/test_integration.py
@@ -585,3 +585,46 @@ def test_completeness_summary_table(model_class):
         f"{model_class.__name__} is not in the summary table. "
         f"Please add it to the summary table."
     )
+
+
+def test_if_models_with_embedding_parameter(model):
+    # Test if the model that have embedding parameters works changing the default
+    # embedding value
+    model_name = model.__class__.__name__
+    # first step is to inspect the models parameters
+    params = inspect.signature(model.__init__).parameters
+
+    # check if there is any mention to emb_size or embedding_dim or emb
+    # or return_features or feature
+
+    parameters_related_with_emb = ["emb_size", "embedding_dim", "emb",
+                                   "return_features", "feature"]
+
+    # check if any of the parameters related to embedding are present in the model
+    if not any(param in params for param in parameters_related_with_emb):
+        pytest.skip(
+            f"{model_name} does not have an embedding parameter."
+        )
+    else:
+        # getting the parameters name
+        emb_param = next((param for param in params if param in parameters_related_with_emb), None)
+
+    # getting the emb parameter and re-initializing the model
+    # changing only this parameters
+    new_value = 128
+
+    model_dict = {name: (args, defaults) for name, args, defaults in models_mandatory_parameters}
+
+    mandatory_parameters, sig_params = model_dict[model_name]
+    signal_params = get_sp(sig_params, mandatory_parameters)
+    # redefining the embedding parameter
+    signal_params[emb_param] = new_value
+
+    model = models_dict[model_name](**signal_params)
+
+    print(f"Model {model_name} has an embedding parameter {emb_param} with value {new_value}.")
+    # assert not error when print the model
+    try:
+        print(model)
+    except Exception as e:
+        pytest.fail(f"Error printing model {model_name}: {e}")