Fix TCN tuple schema unwrap

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.
+        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()

tests/core/test_tcn.py

@@ -163,6 +163,45 @@ 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.
+        """
+        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()