Update (#4920)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-36-188.ap-northeast-1.compute.internal>

python/dgl/mock_sparse/elementwise_op_sp.py

@@ -73,7 +73,7 @@ def sp_add(
             "The shape of sparse matrix A {} and"
             " B {} are expected to match".format(A.shape, B.shape)
         )
-        C = A.adj + B.adj
+        C = (A.adj + B.adj).coalesce()
         return SparseMatrix(C.indices()[0], C.indices()[1], C.values(), C.shape)
     raise RuntimeError(
         "Elementwise addition between {} and {} is not "

python/dgl/mock_sparse/sp_matrix.py

@@ -62,7 +62,7 @@ class SparseMatrix:
         shape: Optional[Tuple[int, int]] = None,
     ):
         if val is None:
-            val = torch.ones(row.shape[0])
+            val = torch.ones(row.shape[0]).to(row.device)
         i = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
         if shape is None:
             self.adj = torch.sparse_coo_tensor(i, val).coalesce()