Fix test_batch_utils sparse.csr

MolecularAI · Aug 13, 2024 · 8bd7d46 · 8bd7d46
1 parent 190dd20
commit 8bd7d46
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_batch_utils.py b/tests/test_batch_utils.py
@@ -119,9 +119,9 @@ def test_combine_csv_batches(tmpdir):
 
 def test_combine_sparse_matrix_batches(tmpdir):
     filename1 = str(tmpdir / "temp.0.npz")
-    sparse.save_npz(filename1, sparse.csr_array([0, 0, 0, 1]), compressed=True)
+    sparse.save_npz(filename1, sparse.csr_array([[0, 0, 0, 1]]), compressed=True)
     filename2 = str(tmpdir / "temp.1.npz")
-    sparse.save_npz(filename2, sparse.csr_array([1, 1, 1, 0]), compressed=True)
+    sparse.save_npz(filename2, sparse.csr_array([[1, 1, 1, 0]]), compressed=True)
     filename = str(tmpdir / "temp.npz")
 
     assert os.path.exists(filename1)