fixed time step

aleatory-science · Nov 24, 2023 · 10daf55 · 10daf55
1 parent ba2682c
commit 10daf55
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/numpyro/contrib/control_flow/scan.py b/numpyro/contrib/control_flow/scan.py
@@ -319,7 +319,7 @@ def body_fn(wrapped_carry, x):
 
                 elif subs_type == "replay":
                     trace = handlers.trace(seeded_fn).get_trace(carry, x)
-                    replay_trace_i = _replay_wrapper(subs_map, trace, i - 1, length)
+                    replay_trace_i = _replay_wrapper(subs_map, trace, i, length)
                     seeded_fn = handlers.replay(seeded_fn, trace=replay_trace_i)
 
             with handlers.trace() as trace: