[Sparse]Fix device check (dmlc#6409)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-24-117.ap-northeast-1.compute.internal>
DominikaJedynak · Mar 12, 2024 · 80c1b58 · 80c1b58
1 parent c31b946
commit 80c1b58
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/dgl_sparse/src/sddmm.cc b/dgl_sparse/src/sddmm.cc
@@ -61,8 +61,9 @@ void _SDDMMSanityCheck(
       mat1.dtype() == mat2.dtype(),
       "SDDMM: the two dense matrices should have the same dtype.");
   TORCH_CHECK(
-      mat1.device() == mat2.device(),
-      "SDDMM: the two dense matrices should on the same device.");
+      mat1.device() == mat2.device() && sparse_mat->device() == mat2.device(),
+      "SDDMM: the two dense matrices and sparse matrix should on the same "
+      "device.");
 }
 
 torch::Tensor SDDMMAutoGrad::forward(