Aligned.

OpenSourceEconomics · Jun 13, 2024 · 7306aed · 7306aed
1 parent 5e2b4e5
commit 7306aed
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/upper_envelope/fues_jax/fues_jax.py b/src/upper_envelope/fues_jax/fues_jax.py
@@ -212,7 +212,9 @@ def fues_jax_unconstrained(
     """
     # Set default value of final grid size to 1.2 times current if not defined
     n_final_wealth_grid = (
-        int(1.2 * (len(policy))) if n_final_wealth_grid is None else n_final_wealth_grid
+        int(1.2 * endog_grid.shape[0])
+        if n_final_wealth_grid is None
+        else n_final_wealth_grid
     )
 
     idx_sort = jnp.argsort(endog_grid)