From 0d62b5ffeceed10da90124dc6f2c8dadbdc914a1 Mon Sep 17 00:00:00 2001
From: Felipe Bonchristiano <felipe.bonchristiano@outlook.com>
Date: Wed, 26 Aug 2026 11:35:15 -0500
Subject: [PATCH 1/2] Fix BIOT channel embedding ignoring emb_size

channel_tokens was nn.Embedding(n_channels, 256), hardcoding 256 and
crashing when added to the emb_size-dim spectral embedding for any
emb_size != 256. Use emb_size.
---
 pyhealth/models/biot.py | 15 +++++++++++++--
 tests/core/test_biot.py | 26 ++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/pyhealth/models/biot.py b/pyhealth/models/biot.py
index e038e092b..4494ba449 100644
--- a/pyhealth/models/biot.py
+++ b/pyhealth/models/biot.py
@@ -129,6 +129,17 @@ def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
 
 
 class BIOTEncoder(nn.Module):
+    """Encoder for multichannel biosignals.
+
+    Examples:
+        >>> encoder = BIOTEncoder(
+        ...     emb_size=128, heads=8, depth=2, n_channels=18
+        ... )
+        >>> signal = torch.randn(2, 18, 2000)
+        >>> encoder(signal).shape
+        torch.Size([2, 128])
+    """
+
     def __init__(
         self,
         emb_size=256,
@@ -158,7 +169,7 @@ def __init__(
         self.positional_encoding = PositionalEncoding(emb_size)
 
         # channel token, N_channels >= your actual channels
-        self.channel_tokens = nn.Embedding(n_channels, 256)
+        self.channel_tokens = nn.Embedding(n_channels, emb_size)
         self.index = nn.Parameter(
             torch.LongTensor(range(n_channels)), requires_grad=False
         )
@@ -372,4 +383,4 @@ def get_embeddings(self, **kwargs: Any) -> Dict[str, torch.Tensor]:
     print(f"✓ BIOTClassifier forward pass:")
     print(f"  Logits shape: {logits.shape}")
     
-    print("\n✓ All tests passed!")
\ No newline at end of file
+    print("\n✓ All tests passed!")
diff --git a/tests/core/test_biot.py b/tests/core/test_biot.py
index d51eecb41..32836dcb3 100644
--- a/tests/core/test_biot.py
+++ b/tests/core/test_biot.py
@@ -176,6 +176,32 @@ def test_model_different_n_classes(self):
 
         self.assertEqual(ret["logit"].shape[1], 1)
 
+    def test_model_non_default_emb_size(self):
+        """BIOT must honor a non-default emb_size.
+
+        Regression test: the channel-token embedding dimension was hardcoded
+        to 256, so it could not be added to the emb_size-dimensional spectral
+        embedding when emb_size != 256, crashing the forward pass.
+        """
+        model = BIOT(
+            dataset=self.dataset,
+            emb_size=128,
+            heads=8,
+            depth=2,
+            n_fft=200,
+            hop_length=100,
+            n_channels=18,
+        )
+        self.assertEqual(model.biot.biot.channel_tokens.weight.shape[1], 128)
+
+        train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=False)
+        data_batch = next(iter(train_loader))
+        ret = model(**data_batch)
+        ret["loss"].backward()
+
+        expected_size = self.dataset.output_processors["label"].size()
+        self.assertEqual(ret["logit"].shape[1], expected_size)
+
     def test_model(self):
         """Test BIOT"""
         model_small = BIOT(

From 1198321de8438fe03101a2e8ed955a9785f11322 Mon Sep 17 00:00:00 2001
From: Felipe Bonchristiano <felipe.bonchristiano@outlook.com>
Date: Fri, 28 Aug 2026 17:37:03 -0500
Subject: [PATCH 2/2] Fix BIOTEncoder usage example

---
 pyhealth/models/biot.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyhealth/models/biot.py b/pyhealth/models/biot.py
index 4494ba449..d3b4b00f0 100644
--- a/pyhealth/models/biot.py
+++ b/pyhealth/models/biot.py
@@ -132,6 +132,7 @@ class BIOTEncoder(nn.Module):
     """Encoder for multichannel biosignals.
 
     Examples:
+        >>> _ = _get_linear_attention_transformer()
         >>> encoder = BIOTEncoder(
         ...     emb_size=128, heads=8, depth=2, n_channels=18
         ... )