Fix BIOT channel embedding ignoring emb_size

pyhealth/models/biot.py

@@ -129,6 +129,18 @@ def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
 
 
 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
+        ... )
+        >>> signal = torch.randn(2, 18, 2000)
+        >>> encoder(signal).shape
+        torch.Size([2, 128])
+    """
+
     def __init__(
         self,
         emb_size=256,
@@ -158,7 +170,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 +384,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!")
+    print("\n✓ All tests passed!")

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(