From d6e0c3c372eb442a877eaafc0f0296b96e7c05e4 Mon Sep 17 00:00:00 2001
From: Felipe Bonchristiano <felipe.bonchristiano@outlook.com>
Date: Thu, 10 Sep 2026 15:07:15 -0500
Subject: [PATCH 1/2] Fix CNN training on singleton vector batches

---
 pyhealth/models/cnn.py | 23 ++++++++++++++++++++---
 tests/core/test_cnn.py | 13 ++++++++++---
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/pyhealth/models/cnn.py b/pyhealth/models/cnn.py
index 29437a722..94e410fd3 100644
--- a/pyhealth/models/cnn.py
+++ b/pyhealth/models/cnn.py
@@ -61,6 +61,23 @@ def __init__(self, in_channels: int, out_channels: int, spatial_dim: int):
             )
         self.relu = nn.ReLU()
 
+    @staticmethod
+    def _apply_layers(layers: nn.Sequential, x: torch.Tensor) -> torch.Tensor:
+        for layer in layers:
+            if (
+                isinstance(layer, nn.BatchNorm1d)
+                and layer.training
+                and x.size(0) * x.size(2) == 1
+            ):
+                # Use running statistics without changing the layer's mode.
+                x = nn.functional.batch_norm(
+                    x, layer.running_mean, layer.running_var,
+                    layer.weight, layer.bias, training=False, eps=layer.eps,
+                )
+            else:
+                x = layer(x)
+        return x
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward propagation.
 
@@ -71,10 +88,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             output tensor of shape [batch size, out_channels, *].
         """
         residual = x
-        out = self.conv1(x)
-        out = self.conv2(out)
+        out = self._apply_layers(self.conv1, x)
+        out = self._apply_layers(self.conv2, out)
         if self.downsample is not None:
-            residual = self.downsample(x)
+            residual = self._apply_layers(self.downsample, x)
         out += residual
         out = self.relu(out)
         return out
diff --git a/tests/core/test_cnn.py b/tests/core/test_cnn.py
index e05b18f7a..54584e318 100644
--- a/tests/core/test_cnn.py
+++ b/tests/core/test_cnn.py
@@ -277,15 +277,22 @@ def test_model_with_multihot_and_1d_tensor_inputs(self):
         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)
+        train_loader = get_dataloader(dataset, batch_size=1, 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["y_prob"].shape[0], 1)
+        self.assertEqual(ret["logit"].shape[0], 1)
         self.assertEqual(ret["loss"].dim(), 0)
+        self.assertTrue(torch.isfinite(ret["loss"]))
+        for module in model.modules():
+            if isinstance(module, torch.nn.BatchNorm1d):
+                self.assertTrue(module.training)
+                self.assertEqual(module.num_batches_tracked.item(), 0)
+                self.assertIsNotNone(module.weight.grad)
+                self.assertTrue(torch.isfinite(module.weight.grad).all())
 
 
 if __name__ == "__main__":

From fdd527d4247f105b3a87dd2c8ccb924dffed7deb Mon Sep 17 00:00:00 2001
From: Felipe Bonchristiano <felipe.bonchristiano@outlook.com>
Date: Thu, 10 Sep 2026 15:16:29 -0500
Subject: [PATCH 2/2] Document CNN singleton batch normalization behavior

---
 docs/api/models/pyhealth.models.CNN.rst | 3 +++
 examples/cnn_mimic4.ipynb               | 4 +++-
 pyhealth/models/cnn.py                  | 7 +++++++
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/docs/api/models/pyhealth.models.CNN.rst b/docs/api/models/pyhealth.models.CNN.rst
index ed5d43ca9..638d830a3 100644
--- a/docs/api/models/pyhealth.models.CNN.rst
+++ b/docs/api/models/pyhealth.models.CNN.rst
@@ -3,6 +3,9 @@
 
 The separate callable CNNLayer and the complete CNN model.
 
+For 1D inputs with one sample and one sequence position, batch normalization
+uses running statistics while preserving gradients and the layers' training state.
+
 .. autoclass:: pyhealth.models.CNNLayer
     :members:
     :undoc-members:
diff --git a/examples/cnn_mimic4.ipynb b/examples/cnn_mimic4.ipynb
index da3868209..44b7eb4fc 100644
--- a/examples/cnn_mimic4.ipynb
+++ b/examples/cnn_mimic4.ipynb
@@ -251,7 +251,9 @@
    "metadata": {},
    "source": [
     "# 5. Instantiate CNN Model\n",
-    "Create the PyHealth CNN with custom hyperparameters and inspect the parameter footprint prior to optimisation."
+    "Create the PyHealth CNN with custom hyperparameters and inspect the parameter footprint prior to optimisation.\n",
+    "\n",
+    "If a final batch contains one patient with a length-1 feature, CNN uses BatchNorm's running statistics for that feature while keeping gradients enabled."
    ]
   },
   {
diff --git a/pyhealth/models/cnn.py b/pyhealth/models/cnn.py
index 94e410fd3..d362d595d 100644
--- a/pyhealth/models/cnn.py
+++ b/pyhealth/models/cnn.py
@@ -31,6 +31,13 @@ class CNNBlock(nn.Module):
     Args:
         in_channels: number of input channels.
         out_channels: number of output channels.
+
+    Example:
+        >>> import torch
+        >>> from pyhealth.models.cnn import CNNBlock
+        >>> block = CNNBlock(4, 8, spatial_dim=1)
+        >>> block(torch.randn(1, 4, 1)).shape
+        torch.Size([1, 8, 1])
     """
 
     def __init__(self, in_channels: int, out_channels: int, spatial_dim: int):