edge_index_cpu = edge_index.cpu()
    for edge_idx in edge_argsort.tolist():
        source = edge_index_cpu[0, edge_idx].item()
        if source not in nodes_remaining:
            continue

        target = edge_index_cpu[1, edge_idx].item()
        if target not in nodes_remaining:
            continue

        new_edge_indices.append(edge_idx)

        cluster[source] = i
        nodes_remaining.remove(source)

        if source != target:
            cluster[target] = i
            nodes_remaining.remove(target)

        i += 1

    # The remaining nodes are simply kept.
    for node_idx in nodes_remaining:
        cluster[node_idx] = i
        i += 1
    cluster = cluster.to(x.device)