cuda tests

test/utils/test_consecutive.py

+import pytest
 import torch
 from torch_cluster.functions.utils.consecutive import consecutive
 
@@ -10,5 +11,18 @@ def test_consecutive():
     assert consecutive(vec).tolist() == [0, 2, 1, 1, 2]
 
     vec = torch.LongTensor([0, 3, 2, 2, 3])
-    assert consecutive(vec, return_unique=True)[0].tolist() == [0, 2, 1, 1, 2]
-    assert consecutive(vec, return_unique=True)[1].tolist() == [0, 2, 3]
+    assert consecutive(vec, True)[0].tolist() == [0, 2, 1, 1, 2]
+    assert consecutive(vec, True)[1].tolist() == [0, 2, 3]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
+def test_consecutive_gpu():  # pragma: no cover
+    vec = torch.cuda.LongTensor([0, 2, 3])
+    assert consecutive(vec).cpu().tolist() == [0, 1, 2]
+
+    vec = torch.cuda.LongTensor([0, 3, 2, 2, 3])
+    assert consecutive(vec).cpu().tolist() == [0, 2, 1, 1, 2]
+
+    vec = torch.cuda.LongTensor([0, 3, 2, 2, 3])
+    assert consecutive(vec, True)[0].cpu().tolist() == [0, 2, 1, 1, 2]
+    assert consecutive(vec, True)[1].cpu().tolist() == [0, 2, 3]

test/utils/test_degree.py

+import pytest
 import torch
 from torch_cluster.functions.utils.degree import node_degree
 
 
-def test_node_degree():
+def test_node_degree_cpu():
     row = torch.LongTensor([0, 1, 1, 0, 0, 3, 0])
     expected_degree = [4, 2, 0, 1]
 
     degree = node_degree(row, 4)
-    assert degree.type() == torch.FloatTensor().type()
+    assert degree.type() == torch.LongTensor().type()
     assert degree.tolist() == expected_degree
 
-    degree = node_degree(row, 4, out=torch.LongTensor())
-    assert degree.type() == torch.LongTensor().type()
+    degree = node_degree(row, 4, out=torch.FloatTensor())
+    assert degree.type() == torch.FloatTensor().type()
     assert degree.tolist() == expected_degree
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
+def test_node_degree_gpu():  # pragma: no cover
+    row = torch.cuda.LongTensor([0, 1, 1, 0, 0, 3, 0])
+    expected_degree = [4, 2, 0, 1]
+
+    degree = node_degree(row, 4)
+    assert degree.type() == torch.cuda.LongTensor().type()
+    assert degree.cpu().tolist() == expected_degree
+
+    degree = node_degree(row, 4, out=torch.cuda.FloatTensor())
+    assert degree.type() == torch.cuda.FloatTensor().type()
+    assert degree.cpu().tolist() == expected_degree

torch_cluster/functions/utils/degree.py

@@ -2,6 +2,7 @@ import torch
 
 
 def node_degree(row, num_nodes, out=None):
+    out = row.new() if out is None else out
     zero = torch.zeros(num_nodes, out=out)
     one = torch.ones(row.size(0), out=zero.new())
     return zero.scatter_add_(0, row, one)