Update src/numpyro_ext/optim.py

Co-authored-by: Adrian Price-Whelan <[email protected]>
dfm · Nov 20, 2022 · f9f4ad9 · f9f4ad9
1 parent b0f10c2
commit f9f4ad9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/numpyro_ext/optim.py b/src/numpyro_ext/optim.py
@@ -84,7 +84,7 @@ def run(rng_key, *args, **kwargs):
         state = svi.init(init_key, *args, **kwargs)
         for _ in range(num_steps):
             state, _ = svi.update(state, *args, **kwargs)
-        info = state.optim_state[1].state
+        info = getattr(state.optim_state[1], 'state', None)
         params = svi.get_params(state)
         sample = guide.sample_posterior(sample_key, params)
         if include_deterministics: