improving use of cpu device

cai4cai · Nov 10, 2022 · b1a5fbb · b1a5fbb
1 parent a127801
commit b1a5fbb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/tests/test_sparse_matmul.py b/tests/test_sparse_matmul.py
@@ -96,7 +96,7 @@ def test_matmul_gradient_csr(self):
 
 
 class SparseMatMulTestCUDA(SparseMatMulTest):
-    """Override superclass setUp to run on CPU"""
+    """Override superclass setUp to run on GPU"""
 
     def setUp(self) -> None:
         if not torch.cuda.is_available():

diff --git a/tests/test_sparse_solve.py b/tests/test_sparse_solve.py
@@ -226,8 +226,8 @@ def setUp(self) -> None:
         self.A_shape = (16, 16)  # square matrix
         self.B_shape = (self.A_shape[1], 10)
         self.A_nnz = 32
-        self.A_idx = gencoordinates_square_tri(self.A_shape[0], self.A_nnz)
-        self.A_val = torch.randn(self.A_nnz, dtype=torch.float64, device="cuda")
+        self.A_idx = gencoordinates_square_tri(self.A_shape[0], self.A_nnz, device=self.device)
+        self.A_val = torch.randn(self.A_nnz, dtype=torch.float64, device=self.device)
 
         self.As_coo_triu = torch.sparse_coo_tensor(self.A_idx, self.A_val, self.A_shape, requires_grad=True).coalesce()
         self.As_coo_tril = self.As_coo_triu.t()
@@ -237,12 +237,12 @@ def setUp(self) -> None:
         self.Ad_triu = self.As_coo_triu.to_dense()
         self.Ad_tril = self.As_coo_tril.to_dense()
 
-        self.Bd = torch.randn(16, 4, dtype=torch.float64, device="cuda")
+        self.Bd = torch.randn(16, 4, dtype=torch.float64, device=self.device)
         self.solve = sparse_triangular_solve
 
 
 class SparseUnitTriangularSolveTestCUDA(SparseUnitTriangularSolveTest):
-    """Override superclass setUp to run on CPU"""
+    """Override superclass setUp to run on GPU"""
 
     def setUp(self) -> None:
         if not torch.cuda.is_available():