test_consecutive.py

import pytest
import torch
from torch_cluster.functions.utils.consecutive import consecutive


def test_consecutive_cpu():
    x = torch.LongTensor([0, 3, 2, 2, 3])
    assert consecutive(x).tolist() == [0, 2, 1, 1, 2]


@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
def test_consecutive_gpu():  # pragma: no cover
    x = torch.cuda.LongTensor([0, 3, 2, 2, 3])
    assert consecutive(x).cpu().tolist() == [0, 2, 1, 1, 2]