Fix index error on cuda

broadinstitute · Aug 30, 2023 · 8be31d2 · 8be31d2
1 parent 89abac9
commit 8be31d2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cellbender/remove_background/posterior.py b/cellbender/remove_background/posterior.py
@@ -535,7 +535,8 @@ def _get_cell_noise_count_posterior_coo(
             m_i = self.index_converter.get_m_indices(cell_inds=bcs_i, gene_inds=genes_i)
 
             nonzero_noise_offset_dict.update(
-                dict(zip(m_i[nonzero_offset_inds], nonzero_noise_count_offsets))
+                dict(zip(m_i[nonzero_offset_inds.detach().cpu()].tolist(),
+                         nonzero_noise_count_offsets.detach().cpu().tolist()))
             )
             c_offset.append(noise_count_offset_NG[bcs_i_chunk, genes_i_analyzed].detach().cpu())