Small fix to mjx tutorial.

PiperOrigin-RevId: 616235636 Change-Id: I38675c032bdca8f26ab0550896928d097ee86653
2024-03-15 13:59:29 -07:00
parent 74d3d8cb4e
commit e95159b4f6
1 changed files with 1 additions and 1 deletions
@@ -367,7 +367,7 @@
    "rng = jax.random.split(rng, 4096)\n",
    "batch = jax.vmap(lambda rng: mjx_data.replace(qpos=jax.random.uniform(rng, (1,))))(rng)\n",
    "\n",
-    "jit_step = jax.vmap(mjx.step, in_axes=(None, 0))\n",
+    "jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))\n",
    "batch = jit_step(mjx_model, batch)\n",
    "\n",
    "print(batch.qpos)"