Changes got clobbered

cwindolf · Dec 6, 2024 · 2b07d8d · 2b07d8d
1 parent 22c5f33
commit 2b07d8d
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/src/dartsort/cluster/gaussian_mixture.py b/src/dartsort/cluster/gaussian_mixture.py
@@ -1603,7 +1603,6 @@ def refit_group(
             if weights is not None:
                 my_weights = weights[i]
                 (keep,) = my_weights.nonzero(as_tuple=True)
-                keep = keep.cpu()
                 my_weights = my_weights[keep]
             u = self.fit_unit(
                 unit_id=k,
@@ -1679,8 +1678,7 @@ def unit_group_criterion(
                 _, sll = self.unit_log_likelihoods(unit=u, spikes=spikes_core)
                 if sll is not None:
                     subunit_logliks[i] = sll
-            (keep,) = subunit_logliks.isfinite().any(dim=0).nonzero(as_tuple=True)
-            keep = keep.cpu()
+            (keep,) = subunit_logliks.isfinite().any(dim=0).cpu().nonzero(as_tuple=True)
             subunit_log_props = (
                 F.softmax(subunit_logliks[:, keep], dim=0).mean(1).log_()
             )