From 9e99500fd9e77731c5ce4e0638753bb1b2add7cc Mon Sep 17 00:00:00 2001
From: Felipe Bonchristiano <felipe.bonchristiano@outlook.com>
Date: Wed, 26 Aug 2026 11:32:06 -0500
Subject: [PATCH 1/2] Fix TCN crash on tuple-schema features
 (StageNetProcessor)

TCN passed raw kwargs (with tuple features) to the embedding model,
crashing on StageNetProcessor's (time, value) tuples. Unwrap the 'value'
tensor via the processor schema first, like the sibling sequence models.
---
 pyhealth/models/tcn.py | 23 ++++++++++++++++++++++-
 tests/core/test_tcn.py | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/pyhealth/models/tcn.py b/pyhealth/models/tcn.py
index 3569ea70b..f92c4cda3 100644
--- a/pyhealth/models/tcn.py
+++ b/pyhealth/models/tcn.py
@@ -306,7 +306,28 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
                 - embed (optional): a tensor representing the patient embeddings if requested.
         """
         patient_emb = []
-        embedded = self.embedding_model(kwargs)
+
+        # Tuple-schema features (e.g. StageNetProcessor emits (time, value))
+        # arrive as a tuple; extract the "value" (and optional "mask") tensor
+        # for the embedding model, mirroring the sibling sequence models.
+        inputs = {}
+        masks = {}
+        for feature_key in self.feature_keys:
+            feature = kwargs[feature_key]
+            if isinstance(feature, torch.Tensor):
+                feature = (feature,)
+            schema = self.dataset.input_processors[feature_key].schema()
+            value = feature[schema.index("value")] if "value" in schema else None
+            mask = feature[schema.index("mask")] if "mask" in schema else None
+            if value is None:
+                raise ValueError(
+                    f"Feature '{feature_key}' must contain 'value' in the schema."
+                )
+            inputs[feature_key] = value
+            if mask is not None:
+                masks[feature_key] = mask
+
+        embedded = self.embedding_model(inputs, masks=masks)
         for feature_key in self.feature_keys:
             x = embedded[feature_key]
             mask = (x.sum(dim=-1) != 0).int()
diff --git a/tests/core/test_tcn.py b/tests/core/test_tcn.py
index a8eee4ea4..bfb1e7558 100644
--- a/tests/core/test_tcn.py
+++ b/tests/core/test_tcn.py
@@ -163,6 +163,46 @@ def test_num_channels_as_list(self):
         self.assertIn("loss", ret)
         self.assertIn("y_prob", ret)
 
+    def test_model_with_stagenet_tuple_feature(self):
+        """TCN must handle tuple-schema features (StageNetProcessor).
+
+        Regression test: StageNetProcessor emits a (time, value) tuple per
+        feature. TCN previously passed the raw tuple to the embedding model
+        and crashed; it must unwrap the "value" tensor like the sibling
+        sequence models do.
+        """
+        samples = [
+            {
+                "patient_id": "patient-0",
+                "visit_id": "visit-0",
+                "codes": ([0.0, 2.0, 1.3], ["c1", "c2", "c3"]),
+                "conditions": ["cond-33", "cond-86"],
+                "label": 0,
+            },
+            {
+                "patient_id": "patient-0",
+                "visit_id": "visit-1",
+                "codes": ([0.0, 2.0], ["c1", "c4"]),
+                "conditions": ["cond-33"],
+                "label": 1,
+            },
+        ]
+        dataset = create_sample_dataset(
+            samples=samples,
+            input_schema={"codes": "stagenet", "conditions": "sequence"},
+            output_schema={"label": "binary"},
+            dataset_name="test_stagenet",
+        )
+        model = TCN(dataset=dataset)
+        data_batch = next(iter(get_dataloader(dataset, batch_size=2, shuffle=False)))
+
+        ret = model(**data_batch)
+        ret["loss"].backward()
+
+        self.assertEqual(ret["y_prob"].shape[0], 2)
+        self.assertEqual(ret["logit"].shape[0], 2)
+        self.assertEqual(ret["loss"].dim(), 0)
+
 
 if __name__ == "__main__":
     unittest.main()

From 3951158399723d122cb713e3100132039392c780 Mon Sep 17 00:00:00 2001
From: Felipe Bonchristiano <felipe.bonchristiano@outlook.com>
Date: Thu, 27 Aug 2026 11:29:09 -0500
Subject: [PATCH 2/2] tcn fix comments update

---
 pyhealth/models/tcn.py | 2 +-
 tests/core/test_tcn.py | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/pyhealth/models/tcn.py b/pyhealth/models/tcn.py
index f92c4cda3..da0ba47fe 100644
--- a/pyhealth/models/tcn.py
+++ b/pyhealth/models/tcn.py
@@ -309,7 +309,7 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
 
         # Tuple-schema features (e.g. StageNetProcessor emits (time, value))
         # arrive as a tuple; extract the "value" (and optional "mask") tensor
-        # for the embedding model, mirroring the sibling sequence models.
+        # for the embedding model.
         inputs = {}
         masks = {}
         for feature_key in self.feature_keys:
diff --git a/tests/core/test_tcn.py b/tests/core/test_tcn.py
index bfb1e7558..e49dd1006 100644
--- a/tests/core/test_tcn.py
+++ b/tests/core/test_tcn.py
@@ -168,8 +168,7 @@ def test_model_with_stagenet_tuple_feature(self):
 
         Regression test: StageNetProcessor emits a (time, value) tuple per
         feature. TCN previously passed the raw tuple to the embedding model
-        and crashed; it must unwrap the "value" tensor like the sibling
-        sequence models do.
+        and crashed; it must unwrap the "value" tensor.
         """
         samples = [
             {