fixed 'index derivative is not defined' message

test/test_spmm.py

+from itertools import product
+
+import pytest
 import torch
 from torch_sparse import spmm
 
+from .utils import dtypes, devices, tensor
+
 
-def test_spmm():
-    row = torch.tensor([0, 0, 1, 2, 2])
-    col = torch.tensor([0, 2, 1, 0, 1])
+@pytest.mark.parametrize('dtype,device', product(dtypes, devices))
+def test_spmm(dtype, device):
+    row = torch.tensor([0, 0, 1, 2, 2], device=device)
+    col = torch.tensor([0, 2, 1, 0, 1], device=device)
     index = torch.stack([row, col], dim=0)
-    value = torch.tensor([1, 2, 4, 1, 3])
+    value = tensor([1, 2, 4, 1, 3], dtype, device)
+    x = tensor([[1, 4], [2, 5], [3, 6]], dtype, device)
 
-    matrix = torch.tensor([[1, 4], [2, 5], [3, 6]])
-    out = spmm(index, value, 3, matrix)
+    out = spmm(index, value, 3, x)
     assert out.tolist() == [[7, 16], [8, 20], [7, 19]]

test/test_spspmm_spmm.py

+from itertools import product
+
+import pytest
+import torch
+from torch_sparse import spspmm, spmm
+
+from .utils import dtypes, devices, tensor
+
+
+@pytest.mark.parametrize('dtype,device', product(dtypes, devices))
+def test_spmm_spspmm(dtype, device):
+    row = torch.tensor([0, 0, 1, 2, 2], device=device)
+    col = torch.tensor([0, 2, 1, 0, 1], device=device)
+    index = torch.stack([row, col], dim=0)
+    value = tensor([1, 2, 4, 1, 3], dtype, device)
+    x = tensor([[1, 4], [2, 5], [3, 6]], dtype, device)
+
+    value = value.requires_grad_(True)
+
+    out_index, out_value = spspmm(index, value, index, value, 3, 3, 3)
+    out = spmm(out_index, out_value, 3, x)
+    assert out.size() == (3, 2)

torch_sparse/spspmm.py

@@ -23,7 +23,8 @@ def spspmm(indexA, valueA, indexB, valueB, m, k, n):
 
     :rtype: (:class:`LongTensor`, :class:`Tensor`)
     """
-    return SpSpMM.apply(indexA, valueA, indexB, valueB, m, k, n)
+    index, value = SpSpMM.apply(indexA, valueA, indexB, valueB, m, k, n)
+    return index.detach(), value
 
 
 class SpSpMM(torch.autograd.Function):