Keep EEGPT channel IDs out of the state dict so pretrained weights load on any montage by bruAristimunha · Pull Request #1195 · braindecode/braindecode

bruAristimunha · 2026-09-30T07:57:17Z

Summary

EEGPT.chans_id is derived from chs_info when the model is built, but it is registered as a persistent buffer. It is therefore saved in checkpoints and copied back on load, which causes two failures with the released braindecode/eegpt-pretrained weights:

The checkpoint does not load unless the model has exactly the same 62 channels. The checkpoint stores 62 IDs, so any other montage fails with size mismatch for chans_id. This includes EEGPT.from_pretrained("braindecode/eegpt-pretrained") with no arguments: the default channel projection (conv1d_constraint) gives 19 IDs.
With the same channel count, loading silently gives the model the wrong channel IDs. The checkpoint's 62 IDs replace the model's own, so a montage with the same channels in a different order looks up the wrong channel embedding for every channel, without any warning.

Reproduced with the released checkpoint (revision e41cb3ae):

`from_pretrained(...)`	master	this PR
no arguments (default channel projection)	size mismatch `[1, 62]` vs `[1, 19]`	loads
`chan_proj_type="none"`, saved 62 channels	loads	loads, same logits and features
`chan_proj_type="none"`, 9 channels	size mismatch `[1, 62]` vs `[1, 9]`	loads, IDs of the 9 channels
`chan_proj_type="none"`, the 62 channels reversed	loads, but with the checkpoint's IDs `[0, 1, 2, …]`	loads, IDs `[61, 60, 59, …]`

Changes

Register chans_id with persistent=False, so it is no longer written to state dicts.
EEGPT._load_from_state_dict drops a chans_id key if present, so checkpoints saved before this change, including the released weights, load under strict=True.
The docstring notes that the IDs are rebuilt from the montage and are not stored.

Testing

New tests in test_models.py:
- test_eegpt_chans_id_is_not_saved;
- test_eegpt_checkpoint_keeps_the_model_channel_ids, which loads an old-style state dict (with chans_id) into a model with fewer channels and into one with the same channels in another order. It checks that the model keeps its own IDs and that the pretrained channel-embedding table still loads;
- test_eegpt_checkpoint_loads_with_channel_projection, where the only missing keys are those of chan_proj.
All 4 cases fail on master, with the size mismatch or the replaced IDs.
Released checkpoint: for the saved 62-channel montage with chan_proj_type="none" and a seeded head, logits and encoder features are bit-identical to master. The classifier probe is not in the checkpoint.
pytest test_models.py test_integration.py test_interpolated.py test_huggingface.py test_return_features.py -k eegpt: 66 passed, 7 skipped. The skips predate this PR: TorchScript, plus activation, dropout and embedding checks.
New or changed behavior is covered by relevant regression tests
Style checks recorded: pre-commit run --files braindecode/models/eegpt.py test/unit_tests/models/test_models.py docs/whats_new.rst passed
docs/whats_new.rst updated

Notes for reviewers

Checkpoints saved after this change no longer contain chans_id. Older braindecode versions will report it as missing when loading them; with strict=False (the from_pretrained default) they rebuild it from chs_info.
A model built without chs_info uses sequential IDs, as before. The checkpoint no longer overrides them; before, it failed to load because the shapes differed ((n_chans,) vs (1, 62)).
Found while fine-tuning EEGPT on several montages in NeuralBench, which currently works around it by subclassing EEGPT.load_state_dict.

chans_id is derived from chs_info at construction but was a persistent buffer, so it was saved in checkpoints and loaded back: - the released braindecode/eegpt-pretrained weights (62 IDs) failed to load with a size mismatch on any other montage, and with the default channel projection (19 IDs), i.e. EEGPT.from_pretrained() with no arguments; - on a montage with the same channel count, loading silently replaced the model's channel IDs with the checkpoint's. Register chans_id as non-persistent and drop the key when loading older checkpoints that still carry it.

braindecode#1194 adds its entry at the top of the same list; keeping the two apart lets them merge in either order without a conflict.

Copilot

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

Review effort: Lite
Findings: 3

Open (3)

Consider matching torch.nn.Module._load_from_state_dict’s full explicit signature (e.g.,… · New
The test channel names use uppercase CZ/PZ. If prepare_chan_ids ever becomes case-sensitive… · New
The test channel names use uppercase CZ/PZ. If prepare_chan_ids ever becomes case-sensitive… · New

What changed in this PR

Ensure EEGPT pretrained checkpoints load across different channel montages by preventing montage-derived chans_id from being persisted/overridden during checkpoint load.

Changes:

Make EEGPT.chans_id a non-persistent buffer so it is excluded from state_dict.
Drop legacy chans_id entries during EEGPT checkpoint loading to keep a model’s montage-derived IDs.
Add regression tests and a release-note entry documenting the new behavior.

File	Description
`braindecode/models/eegpt.py`	Marks `chans_id` as non-persistent and ignores legacy `chans_id` keys on load to avoid montage mismatch/override.
`test/unit_tests/models/test_models.py`	Adds regression tests ensuring `chans_id` isn’t saved and legacy checkpoints don’t override montage-derived IDs.
`docs/whats_new.rst`	Documents the checkpoint-loading behavior change for EEGPT channel IDs.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        # Checkpoints saved before ``chans_id`` became non-persistent (including
+        # the released EEGPT weights) still carry it. Drop it, so the IDs keep
+        # matching this model's montage and a different montage can load.
+        state_dict.pop(prefix + "chans_id", None)
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)


+        ["C3", "CZ", "C4"],  # fewer channels than the checkpoint
+        ["P4", "PZ", "CP2", "C4", "CZ", "C3"],  # same channels, other order


+    model = _eegpt(["C3", "CZ", "C4"])
+    assert "chans_id" not in model.state_dict()
+
+
+@pytest.mark.parametrize(
+    "target_names",
+    [
+        ["C3", "CZ", "C4"],  # fewer channels than the checkpoint
+        ["P4", "PZ", "CP2", "C4", "CZ", "C3"],  # same channels, other order
+    ],
+)
+def test_eegpt_checkpoint_keeps_the_model_channel_ids(target_names):
+    # Checkpoints saved before chans_id became non-persistent, like the released
+    # EEGPT weights, still store it. Loading one must neither fail on another
+    # montage nor replace the IDs of the model's channels by the checkpoint's.
+    source = _eegpt(["C3", "CZ", "C4", "CP2", "PZ", "P4"])


Copilot

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

Review effort: Lite
Findings: 2 · 3

Open (5)

This mutates the caller-provided state_dict in-place, which can be surprising if the same dict is… · New
The all(...) assertion is vacuously true when missing_keys is empty, so this test could pass… · New
The test channel names use uppercase CZ/PZ. If prepare_chan_ids ever becomes case-sensitive…
The test channel names use uppercase CZ/PZ. If prepare_chan_ids ever becomes case-sensitive…
Consider matching torch.nn.Module._load_from_state_dict’s full explicit signature (e.g.,…

+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        # Checkpoints saved before ``chans_id`` became non-persistent (including
+        # the released EEGPT weights) still carry it. Drop it, so the IDs keep
+        # matching this model's montage and a different montage can load.


+
+    incompatible = target.load_state_dict(state_dict, strict=False)
+
+    assert not incompatible.unexpected_keys


codecov · 2026-09-30T09:01:05Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87.72%. Comparing base (017809e) to head (1c143de).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1195   +/-   ##
=======================================
  Coverage   87.71%   87.72%           
=======================================
  Files         149      149           
  Lines       17433    17436    +3     
=======================================
+ Hits        15292    15295    +3     
  Misses       2141     2141

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copilot AI balanced review requested due to automatic review settings September 30, 2026 07:57

bruAristimunha added 2 commits September 30, 2026 09:57

Add whats_new entry for braindecode#1195

6414168

Move the braindecode#1195 whats_new entry to the end of the bug-fix list

1c143de

braindecode#1194 adds its entry at the top of the same list; keeping the two apart lets them merge in either order without a conflict.

Copilot AI reviewed Sep 30, 2026

View reviewed changes

Copilot AI balanced review requested due to automatic review settings September 30, 2026 08:14

Copilot started reviewing on behalf of bruAristimunha September 30, 2026 08:28 View session

Copilot AI reviewed Sep 30, 2026

View reviewed changes

Copilot started reviewing on behalf of bruAristimunha September 30, 2026 09:18 View session

bruAristimunha merged commit ccbbea3 into braindecode:master Oct 4, 2026
14 checks passed

bruAristimunha mentioned this pull request Oct 5, 2026

Pretrained models compatible by default: input-geometry contract, audits of the 19 pretrained classes, stacked plan #1227

Open

Navigation Menu

Repository navigation

Keep EEGPT channel IDs out of the state dict so pretrained weights load on any montage - #1195

Keep EEGPT channel IDs out of the state dict so pretrained weights load on any montage#1195
bruAristimunha merged 3 commits into
braindecode:masterfrom
bruAristimunha:fix/eegpt-chans-id-buffer

bruAristimunha commented Sep 30, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

codecov Bot commented Sep 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		["C3", "CZ", "C4"], # fewer channels than the checkpoint
		["P4", "PZ", "CP2", "C4", "CZ", "C3"], # same channels, other order


		incompatible = target.load_state_dict(state_dict, strict=False)

		assert not incompatible.unexpected_keys

Conversation

bruAristimunha commented Sep 30, 2026

Summary

Changes

Testing

Notes for reviewers

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

Uh oh!

codecov Bot commented Sep 30, 2026

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants