enhance codes

mancusolab · Jul 8, 2024 · b7955f0 · b7955f0
1 parent b05684b
commit b7955f0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/sushie/infer.py b/sushie/infer.py
@@ -469,11 +469,11 @@ def infer_sushie(
         )
         elbo_last = elbo_tracker[o_iter]
         elbo_tracker = jnp.append(elbo_tracker, elbo_cur)
-        elbo_increase = not (
-            elbo_cur < elbo_last and (not jnp.isclose(elbo_cur, elbo_last, atol=1e-8))
+        elbo_increase = elbo_cur >= elbo_last or jnp.isclose(
+            elbo_cur, elbo_last, atol=1e-8
         )
 
-        if (not elbo_increase) or jnp.isnan(elbo_cur):
+        if not elbo_increase:
             log.logger.warning(
                 f"Optimization concludes after {o_iter + 1} iterations."
                 + f" ELBO decreases. Final ELBO score: {elbo_cur}. Return last iteration's results."