From 4e9ec07f184076054c5fcbca67015db94d647f32 Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Wed, 26 Aug 2026 11:21:20 -0500 Subject: [PATCH 1/2] Fix CNN crash on 1D tensor and multi-hot inputs CNN.forward hardcoded a 3-D expectation for spatial_dim=1 features, but MultiHotProcessor and 1D TensorProcessor inputs embed to [batch, embedding_dim] with no sequence axis. Treat these as a length-1 sequence. --- pyhealth/models/cnn.py | 6 ++++++ tests/core/test_cnn.py | 48 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/pyhealth/models/cnn.py b/pyhealth/models/cnn.py index 3d5fe3ce0..4c13e182b 100644 --- a/pyhealth/models/cnn.py +++ b/pyhealth/models/cnn.py @@ -329,6 +329,12 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: x = x.to(self.device) spatial_dim = self.feature_conv_dims[feature_key] + # Some spatial_dim==1 features embed to [batch, embedding_dim] with no + # sequence axis (e.g. MultiHotProcessor, or a TensorProcessor whose + # per-sample value is a 1D vector). Treat these as a length-1 sequence + # so the 1D CNN can run on them. + if spatial_dim == 1 and x.dim() == 2: + x = x.unsqueeze(1) expected_dims = {1: 3, 2: 4, 3: 5}[spatial_dim] if x.dim() != expected_dims: raise ValueError( diff --git a/tests/core/test_cnn.py b/tests/core/test_cnn.py index 4a53e5444..e05b18f7a 100644 --- a/tests/core/test_cnn.py +++ b/tests/core/test_cnn.py @@ -239,6 +239,54 @@ def test_model_with_mixed_inputs(self): self.assertEqual(ret["logit"].shape[0], 2) self.assertEqual(ret["loss"].dim(), 0) + def test_model_with_multihot_and_1d_tensor_inputs(self): + """Test CNN model with non-sequence spatial_dim=1 inputs. + + MultiHotProcessor and a TensorProcessor whose per-sample value is a 1D + vector both embed to [batch, embedding_dim] with no sequence axis. These + are documented as supported input types and must not crash. + """ + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "demographics": ["asian", "non_hispanic"], + "vitals": [1.0, 2.5, 3.0], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "demographics": ["white"], + "vitals": [0.5, 1.0, 2.0], + "label": 0, + }, + ] + + input_schema = {"demographics": "multi_hot", "vitals": "tensor"} + output_schema = {"label": "binary"} + + dataset = create_sample_dataset( + samples=samples, + input_schema=input_schema, + output_schema=output_schema, + dataset_name="test_multihot", + ) + + model = CNN(dataset=dataset) + self.assertEqual(model.feature_conv_dims["demographics"], 1) + self.assertEqual(model.feature_conv_dims["vitals"], 1) + + train_loader = get_dataloader(dataset, batch_size=2, shuffle=False) + data_batch = next(iter(train_loader)) + + 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 7eed6e99c90f6ceaa4e93fd1b3b9f33bee7caaf9 Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Wed, 26 Aug 2026 23:45:07 -0500 Subject: [PATCH 2/2] changed cnn fix comment --- pyhealth/models/cnn.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/pyhealth/models/cnn.py b/pyhealth/models/cnn.py index 4c13e182b..29437a722 100644 --- a/pyhealth/models/cnn.py +++ b/pyhealth/models/cnn.py @@ -329,10 +329,8 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: x = x.to(self.device) spatial_dim = self.feature_conv_dims[feature_key] - # Some spatial_dim==1 features embed to [batch, embedding_dim] with no - # sequence axis (e.g. MultiHotProcessor, or a TensorProcessor whose - # per-sample value is a 1D vector). Treat these as a length-1 sequence - # so the 1D CNN can run on them. + # Treat spatial_dim==1 features (which embed to [batch, embedding_dim] + # with no sequence axis) as a length-1 sequence so the 1D CNN can run on them. if spatial_dim == 1 and x.dim() == 2: x = x.unsqueeze(1) expected_dims = {1: 3, 2: 4, 3: 5}[spatial_dim]