9ec3b8e2099639500735b3074f8a02c304ba1095,geomstats/learning/em_expectation_maximization.py,RiemannianEM,update_means,#RiemannianEM#Any#Any#Any#Any#Any#Any#,92

Before Change

        if(g_index>0):
            self.means[g_index] = barycenter(data, wik[:, g_index], lr_mu, tau_mu, max_iter=max_iter, normed=True).squeeze()
        else:
            self.means = barycenter(data.unsqueeze().expand(N, M, D), wik, lr_mu, tau_mu, max_iter=max_iter, normed=True).squeeze()


    def update_variances(self, z, wik, g_index=-1):

After Change

        if(g_index>0):
            mean.fit(data.data.from_numpy(), weights=wik.data.numpy()[:,g_index])
            self.means[g_index] = torch.from_numpy(mean.estimate_).squeeze()

        else:
            mean.fit(data_gs, weights = wik.data.numpy())

Instances