diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py
index 0ef94281..d0bd1110 100644
--- a/mjx/mujoco/mjx/_src/render_util.py
+++ b/mjx/mujoco/mjx/_src/render_util.py
@@ -110,4 +110,4 @@ def get_depth(
   )
 
   depth = jnp.clip(raw / depth_scale, 0.0, 1.0)
-  return depth.reshape(height, width)
+  return depth.reshape(height, width, 1)
diff --git a/mjx/mujoco/mjx/_src/render_util_test.py b/mjx/mujoco/mjx/_src/render_util_test.py
index 46c74420..7549b0af 100644
--- a/mjx/mujoco/mjx/_src/render_util_test.py
+++ b/mjx/mujoco/mjx/_src/render_util_test.py
@@ -94,7 +94,7 @@ class RenderUtilTest(absltest.TestCase):
           rc, 0, depth_data, 5.0
       )
 
-    self.assertEqual(depth.shape, (height, width))
+    self.assertEqual(depth.shape, (height, width, 1))
 
   def test_get_depth_vmap(self):
     nworld, width, height = 3, 4, 4
@@ -111,7 +111,7 @@ class RenderUtilTest(absltest.TestCase):
           static_argnums=(0, 1, 3),
       )(rc, 0, depth_data, 5.0)
 
-    self.assertEqual(depth.shape, (nworld, height, width))
+    self.assertEqual(depth.shape, (nworld, height, width, 1))
 
 
 if __name__ == '__main__':