added gpu tests

test/test_grid.py

@@ -29,4 +29,20 @@ def test_grid_cluster_cpu(tensor):
 @pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
 @pytest.mark.parametrize('tensor', tensors)
 def test_grid_cluster_gpu(tensor):  # pragma: no cover
-    pass
+    position = Tensor(tensor, [[0, 0], [9, 9], [2, 8], [2, 2], [8, 3]]).cuda()
+    size = torch.cuda.LongTensor([5, 5])
+    expected = torch.LongTensor([0, 3, 1, 0, 2])
+
+    output = grid_cluster(position, size)
+    assert output.cpu().tolist() == expected.tolist()
+
+    output = grid_cluster(position.expand(2, 5, 2), size)
+    assert output.cpu().tolist() == expected.expand(2, 5).tolist()
+
+    expected = torch.LongTensor([0, 1, 3, 2, 4])
+    batch = torch.cuda.LongTensor([0, 0, 1, 1, 1])
+    output = grid_cluster(position, size, batch)
+    assert output.cpu().tolist() == expected.tolist()
+
+    output = grid_cluster(position.expand(2, 5, 2), size, batch.expand(2, 5))
+    assert output.cpu().tolist() == expected.expand(2, 5).tolist()