diff --git a/brax/io/json.py b/brax/io/json.py
index bb78482f..2b90c5f8 100644
--- a/brax/io/json.py
+++ b/brax/io/json.py
@@ -171,7 +171,8 @@ def dumps(sys: System, states: List[State]) -> Text:
     }
 
     if geom['name'] in ('Mesh', 'Box'):
-      vert, face = _get_mesh(sys.mj_model, sys.geom_dataid[id_])
+      # TODO: use sys.geom_dataid.
+      vert, face = _get_mesh(sys.mj_model, sys.mj_model.geom_dataid[id_])
       geom['vert'] = vert
       geom['face'] = face
 
diff --git a/brax/visualizer/visualizer.py b/brax/visualizer/visualizer.py
index 061faa5e..2faa3e17 100644
--- a/brax/visualizer/visualizer.py
+++ b/brax/visualizer/visualizer.py
@@ -20,7 +20,7 @@
 
 from absl import app
 from absl import flags
-from brax.base import Actuator, Convex, Mesh, Plane, Transform
+from brax.base import Actuator
 from brax.generalized import pipeline as generalized
 from brax.io import html
 from brax.io import mjcf