fix compile behavior

pyro-ppl · May 22, 2024 · 8534f5f · 8534f5f
1 parent 5ebb63d
commit 8534f5f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/numpyro/util.py b/numpyro/util.py
@@ -382,7 +382,9 @@ def loop_fn(collection):
 
         if upper == 0:
             # special case, only compiling
-            _body_fn_wrap(0, *vals)
+            val, collection, start_idx, thinning = vals
+            dummy_collection = jax.tree.map(lambda x: x.copy(), collection)
+            _body_fn_wrap(0, val, dummy_collection, start_idx, thinning)
         else:
             with tqdm.trange(upper) as t:
                 for i in t: