use init_to_sample as init_strategy

CDCgov · Sep 20, 2024 · a19f1ba · a19f1ba
1 parent 50af19b
commit a19f1ba
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pyrenew/metaclass.py b/pyrenew/metaclass.py
@@ -10,7 +10,7 @@
 import numpy as np
 import polars as pl
 from jax.typing import ArrayLike
-from numpyro.infer import MCMC, NUTS, Predictive
+from numpyro.infer import MCMC, NUTS, Predictive, init_to_sample
 
 from pyrenew.mcmcutils import plot_posterior, spread_draws
 
@@ -176,6 +176,9 @@ def _init_model(
         if "find_heuristic_step_size" not in nuts_args:
             nuts_args["find_heuristic_step_size"] = True
 
+        if "init_strategy" not in nuts_args:
+            nuts_args["init_strategy"] = init_to_sample
+
         if mcmc_args is None:
             mcmc_args = dict()