prevent squeeze from reducing one-column rows to scalar

xinehc · Dec 16, 2024 · e07636d · e07636d
1 parent bbbe15c
commit e07636d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/melon/melon.py b/src/melon/melon.py
@@ -249,8 +249,8 @@ def run_em(self, max_iterations=1000, epsilon=1e-10):
         ## return assignments
         assignments = []
         for p_read in p_reads.tocsr():
-            p_read = p_read.toarray().squeeze()
-            assignments.append(np.where(p_read == p_read.max())[0].tolist())
+            p_read = p_read.toarray()
+            assignments.append(np.where(p_read == p_read.max())[1].tolist())
 
         ties = defaultdict(set)
         for qseqid, lineage in enumerate(assignments):