From e95159b4f6d48d114b16a8dc13ad26b3e44bc3e2 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Fri, 15 Mar 2024 13:59:29 -0700
Subject: [PATCH] Small fix to mjx tutorial.

PiperOrigin-RevId: 616235636
Change-Id: I38675c032bdca8f26ab0550896928d097ee86653
---
 mjx/tutorial.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb
index 65579cad..f57aee4d 100644
--- a/mjx/tutorial.ipynb
+++ b/mjx/tutorial.ipynb
@@ -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)"