diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py
index 0de0cb4..4a50e8a 100644
--- a/dm_pix/_src/augment.py
+++ b/dm_pix/_src/augment.py
@@ -358,7 +358,10 @@ def pad_to_size(
   top = delta_height // 2
   bottom = max(target_height - (top + height), 0)
 
-  pad_width = ((top, bottom), (left, right), (0, 0))
+  if _channels_last(image, channel_axis):
+    pad_width = ((top, bottom), (left, right), (0, 0))
+  else:
+    pad_width = ((0, 0), (top, bottom), (left, right))
   if batch:
     pad_width = ((0, 0), *pad_width)
 
diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py
index f584a17..8532e94 100644
--- a/dm_pix/_src/augment_test.py
+++ b/dm_pix/_src/augment_test.py
@@ -559,6 +559,39 @@ def test_pad_to_size_when_target_size_smaller_than_original(
     self.assertEqual(output.shape[1], expected_height)
     self.assertEqual(output.shape[2], expected_width)
 
+  def test_pad_to_size_hwc_shape(self):
+    image = jnp.zeros((4, 4, 3))
+    result = augment.pad_to_size(
+        image, target_height=6, target_width=6, channel_axis=-1
+    )
+    self.assertEqual(result.shape, (6, 6, 3))
+
+  def test_pad_to_size_chw_shape(self):
+    image = jnp.zeros((3, 4, 4))
+    result = augment.pad_to_size(
+        image, target_height=6, target_width=6, channel_axis=0
+    )
+    self.assertEqual(result.shape, (3, 6, 6))
+
+  def test_pad_to_size_chw_values(self):
+    image = jnp.ones((1, 1, 1))
+    result = augment.pad_to_size(
+        image, target_height=3, target_width=3, channel_axis=0
+    )
+    self.assertEqual(result.shape, (1, 3, 3))
+    np.testing.assert_array_equal(result[0, 1, 1], 1.0)
+    np.testing.assert_array_equal(result[0, 0, :], 0.0)
+    np.testing.assert_array_equal(result[0, 2, :], 0.0)
+    np.testing.assert_array_equal(result[0, :, 0], 0.0)
+    np.testing.assert_array_equal(result[0, :, 2], 0.0)
+
+  def test_pad_to_size_batch_chw_shape(self):
+    image = jnp.zeros((2, 3, 4, 4))
+    result = augment.pad_to_size(
+        image, target_height=6, target_width=6, channel_axis=1
+    )
+    self.assertEqual(result.shape, (2, 3, 6, 6))
+
 
 if __name__ == "__main__":
   jax.config.update("jax_default_matmul_precision", "float32")