[Sparse] Fix the sddmm backward problem in mock_sparse. (#5070)

3eac464a · czkkkkkk · GitHub · 4aff59f7 · 3eac464a
Unverified Commit 3eac464a authored Dec 26, 2022 by czkkkkkk Committed by GitHub Dec 26, 2022
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

python/dgl/mock_sparse/sddmm.py python/dgl/mock_sparse/sddmm.py +3 -1

No files found.
--- a/python/dgl/mock_sparse/sddmm.py
+++ b/python/dgl/mock_sparse/sddmm.py
@@ -54,7 +54,9 @@ def sddmm(
        "scalar values. "
    )
    # PyTorch's sddmm operator only supports CSR format.
-    res = torch.sparse.sampled_addmm(A.adj.to_sparse_csr(), mat1, mat2)
+    res = torch.sparse.sampled_addmm(
+        A.adj.to_sparse_csr(), mat1, mat2
+    ).to_sparse_coo()
    return SparseMatrix(A.row, A.col, res.values(), A.adj.shape)