From b4694ee331e85d789a51651378a4a05f0367d3bc Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Tue, 19 Mar 2024 14:35:13 -0700
Subject: [PATCH] Fix bug in MJX viewer.

PiperOrigin-RevId: 617297574
Change-Id: I34fb797db47308a8b2f29a6c54f2b59d77f84a32
---
 mjx/mujoco/mjx/_src/io.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index 1958829f..d290a612 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -239,6 +239,7 @@ def _get_contact(
 
   ncon = cx.dist.shape[0]
   c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id]
+  c.dim[:] = 3
 
 
 def get_data(