Merge pull request #93 from xuyuon/fix-nan

Handling nan output from Adam optimizer when generating reference parameters
kazewong · Jul 16, 2024 · fe0e767 · fe0e767
2 parents 94c34db + bd71684
commit fe0e767
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/jimgw/single_event/likelihood.py b/src/jimgw/single_event/likelihood.py
@@ -559,7 +559,7 @@ def y(x):
         rng_key, optimized_positions, summary = optimizer.optimize(
             jax.random.PRNGKey(12094), y, initial_position
         )
-        best_fit = optimized_positions[summary["final_log_prob"].argmin()]
+        best_fit = optimized_positions[jnp.nanargmin(summary["final_log_prob"])]
         return prior.transform(prior.add_name(best_fit))