6437d4e459e0f56e059ef0400c03852ef91b8a32,torch_geometric/nn/pool/consecutive.py,,consecutive_cluster,#Any#,16

Before Change

    key, perm = unique(src)
    max_value = key.size(0)
    dtype = _get_dtype(max_value)
    arg = torch.empty((key[-1] + 1, ), dtype=dtype, device=src.device)
    arg[key] = torch.arange(0, max_value, dtype=dtype, device=src.device)
    out = arg[src.view(-1)]
    out = out.view(size).long()

After Change

def consecutive_cluster(src):
    unique, inv = torch.unique(src, sorted=True, return_inverse=True)
    perm = torch.arange(inv.size(0), dtype=inv.dtype, device=inv.device)
    perm = inv.new_empty(unique.size(0)).scatter_(0, inv, perm)
    return inv, perm

Instances