diff --git a/pyhealth/models/tcn.py b/pyhealth/models/tcn.py index 3569ea70b..da0ba47fe 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. + 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..e49dd1006 100644 --- a/tests/core/test_tcn.py +++ b/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()