restore mean tracking in scan version

pyro-ppl · Jun 18, 2024 · 8ebbaa7 · 8ebbaa7
1 parent 17821d1
commit 8ebbaa7
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/ar2.py b/examples/ar2.py
@@ -64,6 +64,8 @@ def transition(carry, _):
     with numpyro.handlers.condition(data={"y": y[2:]}):
         _, mu = scan(transition, init, timesteps)
 
+    numpyro.deterministic("mu", mu)
+
 
 def ar2_for_loop(y):
     alpha_1 = numpyro.sample("alpha_1", dist.Normal(0, 1))