Fix cuda issues in test design

broadinstitute · Aug 29, 2023 · 1a3ce15 · 1a3ce15
1 parent b98b20b
commit 1a3ce15
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/cellbender/remove_background/tests/test_dataprep.py b/cellbender/remove_background/tests/test_dataprep.py
@@ -75,9 +75,9 @@ def test_dataloader_sorting(simulated_dataset, cuda):
             bcs_i = loader.unsort_inds(bcs_i)
 
             # Add sparse matrix values to lists.
-            barcodes.append(bcs_i)
-            genes.append(genes_i)
-            counts.append(counts_i)
+            barcodes.append(bcs_i.detach().cpu())
+            genes.append(genes_i.detach().cpu())
+            counts.append(counts_i.detach().cpu())
 
             # Increment barcode index counter.
             ind += data.shape[0]  # Same as data_loader.batch_size

diff --git a/cellbender/remove_background/tests/test_sparse_utils.py b/cellbender/remove_background/tests/test_sparse_utils.py
@@ -76,9 +76,9 @@ def test_dense_to_sparse_op_torch(simulated_dataset, cuda):
         bcs_i = data_loader.unsort_inds(bcs_i)
 
         # Add sparse matrix values to lists.
-        barcodes.append(bcs_i)
-        genes.append(genes_i)
-        counts.append(counts_i)
+        barcodes.append(bcs_i.detach().cpu())
+        genes.append(genes_i.detach().cpu())
+        counts.append(counts_i.detach().cpu())
 
         # Increment barcode index counter.
         ind += data.shape[0]  # Same as data_loader.batch_size