Update fit.py

adonath · web-flow · commit c7dc7052e58e · 2025-07-08T13:17:23.000-07:00
diff --git a/gmmx/fit.py b/gmmx/fit.py
@@ -161,5 +161,4 @@ def em_cond(
             body_fun=em_step,
             init_val=(x, gmm, 0, jnp.asarray(jnp.inf), jnp.array(jnp.inf)),
         )
-        result = jax.block_until_ready(result)
         return EMFitterResult(*result, converged=result[2] < self.max_iter)  # type: ignore [misc]

Original file line number	Diff line number	Diff line change
`@@ -161,5 +161,4 @@ def em_cond(`
`161`	`161`	`body_fun=em_step,`
`162`	`162`	`init_val=(x, gmm, 0, jnp.asarray(jnp.inf), jnp.array(jnp.inf)),`
`163`	`163`	`)`
`164`		`- result = jax.block_until_ready(result)`
`165`	`164`	`return EMFitterResult(*result, converged=result[2] < self.max_iter) # type: ignore [misc]`