diff --git a/src/av2/rendering/rasterize.py b/src/av2/rendering/rasterize.py
index 33d9e7f0..c1d4c117 100644
--- a/src/av2/rendering/rasterize.py
+++ b/src/av2/rendering/rasterize.py
@@ -56,13 +56,12 @@ def xyz_to_bev(
     # Otherwise, use the provided intensity.
     intensity: NDArrayByte
     if xyz.shape[-1] == 3:
-        intensity = np.ones_like(xyz.shape[0], np.uint8)
+        cart = xyz.copy()
+        intensity = np.ones_like(xyz[:, 0:1], np.uint8)
     else:
+        cart = xyz[..., :-1].copy()
         intensity = xyz[..., -1].copy().astype(np.uint8)
 
-    # Grab the Cartesian coordinates (xyz).
-    cart = xyz[..., :-1].copy()
-
     # Move the origin to the center of the image.
     cart += np.divide(grid_size_m, 2)
 
diff --git a/tests/rendering/test_rasterize.py b/tests/rendering/test_rasterize.py
new file mode 100644
index 00000000..fd206681
--- /dev/null
+++ b/tests/rendering/test_rasterize.py
@@ -0,0 +1,41 @@
+"""Tests for the rasterize module."""
+
+from typing import Tuple
+
+import numpy as np
+
+from av2.rendering.rasterize import xyz_to_bev
+from av2.utils.typing import NDArrayFloat
+
+
+def _build_dummy_raster_inputs(
+    n: int, d: int
+) -> Tuple[NDArrayFloat, Tuple[float, float, float], Tuple[float, float, float], NDArrayFloat]:
+    """Build dummy inputs for the rasterize function.
+
+    Args:
+        n: Number of points.
+        d: Number of dimensions.
+
+    Returns:
+        (n,d) points, (3,) voxel resolution, (3,) grid size, (n,) cmap values.
+    """
+    xyz = np.ones((n, d))
+    voxel_resolution = (0.1, 0.1, 0.1)
+    grid_size_m = (50.0, 50.0, 10.0)
+    cmap = np.ones_like(xyz[:, 0:1])
+    return xyz, voxel_resolution, grid_size_m, cmap
+
+
+def test_rasterize_Nx3() -> None:
+    """Test the rasterize function with (N,3) input."""
+    n, d = 1000, 3
+    xyz, voxel_resolution, grid_size_m, cmap = _build_dummy_raster_inputs(n, d)
+    xyz_to_bev(xyz, voxel_resolution, grid_size_m, cmap)
+
+
+def test_rasterize_Nx4() -> None:
+    """Test the rasterize function with (N,4) input."""
+    n, d = 1000, 4
+    xyz, voxel_resolution, grid_size_m, cmap = _build_dummy_raster_inputs(n, d)
+    xyz_to_bev(xyz, voxel_resolution, grid_size_m, cmap)