diff --git a/src/liger_kernel/ops/group_norm.py b/src/liger_kernel/ops/group_norm.py
index 865fc337f..8de49f1ca 100644
--- a/src/liger_kernel/ops/group_norm.py
+++ b/src/liger_kernel/ops/group_norm.py
@@ -242,9 +242,11 @@ def group_norm_forward(X, num_channels, num_groups, W, B, eps):
 def group_norm_backward(dY, X, W, B, Mean, RSTD, num_channels, num_groups):
     shape = dY.shape
     batch_size = shape[0]
-    hidden_size = dY.shape[-1]
     channels_per_group = num_channels // num_groups
     dY = dY.view(batch_size, num_groups, -1)
+    # Number of elements per channel, so it has to be measured after the flatten:
+    # the last dimension of an unflattened (N, C, H, W) input is only W
+    hidden_size = dY.shape[-1] // channels_per_group
     DX = torch.empty(
         (batch_size, num_groups, hidden_size * channels_per_group),
         dtype=X.dtype,
diff --git a/test/transformers/test_group_norm.py b/test/transformers/test_group_norm.py
index 3454947f9..010bcea87 100644
--- a/test/transformers/test_group_norm.py
+++ b/test/transformers/test_group_norm.py
@@ -7,26 +7,10 @@
 device = infer_device()
 
 
-@pytest.mark.parametrize(
-    "batch_size, num_channels, num_groups, hidden_size",
-    [
-        (1, 1, 1, 3),  # minimal
-        (1, 32, 32, 4),  # group == channel
-        (16, 32, 1, 4096),  # single group
-        (2, 63, 21, 2163),  # non-aligned hidden
-        (16, 48, 12, 8192),  # large hidden
-    ],
-)
-@pytest.mark.parametrize(
-    "dtype, atol, rtol",
-    [
-        (torch.float32, 1e-4, 1e-4),
-    ],
-)
-def test_liger_group_norm(batch_size, num_channels, num_groups, hidden_size, dtype, atol, rtol):
+def _test_liger_group_norm(shape, num_channels, num_groups, dtype, atol, rtol):
     torch.manual_seed(0)
 
-    _tensor = torch.randn(batch_size, num_channels, hidden_size, dtype=dtype, device=device)
+    _tensor = torch.randn(*shape, dtype=dtype, device=device)
 
     liger_x = _tensor.clone().detach().requires_grad_(True)
     torch_x = _tensor.clone().detach().requires_grad_(True)
@@ -50,3 +34,41 @@ def test_liger_group_norm(batch_size, num_channels, num_groups, hidden_size, dty
     assert torch.allclose(liger_x.grad, torch_x.grad, atol=atol, rtol=rtol)
     assert torch.allclose(liger_ln.bias.grad, torch_ln.bias.grad, atol=atol, rtol=rtol), "Bias grads different"
     assert torch.allclose(liger_ln.weight.grad, torch_ln.weight.grad, atol=atol, rtol=rtol), "Weight grads different"
+
+
+@pytest.mark.parametrize(
+    "batch_size, num_channels, num_groups, hidden_size",
+    [
+        (1, 1, 1, 3),  # minimal
+        (1, 32, 32, 4),  # group == channel
+        (16, 32, 1, 4096),  # single group
+        (2, 63, 21, 2163),  # non-aligned hidden
+        (16, 48, 12, 8192),  # large hidden
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype, atol, rtol",
+    [
+        (torch.float32, 1e-4, 1e-4),
+    ],
+)
+def test_liger_group_norm(batch_size, num_channels, num_groups, hidden_size, dtype, atol, rtol):
+    _test_liger_group_norm((batch_size, num_channels, hidden_size), num_channels, num_groups, dtype, atol, rtol)
+
+
+@pytest.mark.parametrize(
+    "shape, num_channels, num_groups",
+    [
+        ((2, 6, 4, 8), 6, 3),  # convolutional features
+        ((4, 32, 16, 16), 32, 8),  # larger spatial plane
+        ((2, 8, 2, 4, 4), 8, 4),  # volumetric features
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype, atol, rtol",
+    [
+        (torch.float32, 1e-4, 1e-4),
+    ],
+)
+def test_liger_group_norm_spatial_dims(shape, num_channels, num_groups, dtype, atol, rtol):
+    _test_liger_group_norm(shape, num_channels, num_groups, dtype, atol, rtol)