Adapt fit.py

adonath · Dec 18, 2024 · 614f510 · 614f510
1 parent 99578f6
commit 614f510
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gmmx/fit.py b/gmmx/fit.py
@@ -100,7 +100,7 @@ def m_step(
             Updated Gaussian mixture model instance.
         """
         x = jnp.expand_dims(x, axis=(Axis.components, Axis.features_covar))
-        return gmm.update_parameters(x, jnp.exp(log_resp), reg_covar=self.reg_covar)
+        return gmm.from_responsibilities(x, jnp.exp(log_resp), reg_covar=self.reg_covar)
 
     @jax.jit
     def fit(self, x: jax.Array, gmm: GaussianMixtureModelJax) -> EMFitterResult: