Thank you, topk can do the work. But I need the topk for each point the data. topk may end up with some overlap. A point belonging to more than one. For example, i have a 12936x4096 tensor. I would like to get pairs of indices that i can merge/fusion (if k=2). The distance between the two pair should be minimum.  I"m currently doing something like that
x = torch.randn(12936, 4096)
y = x
m = len(x)
dists = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, m) + \
        torch.pow(y, 2).sum(dim=1, keepdim=True).expand(m, m).t()
pairs = []

How do I get the pair of index that can be merged?
Thank you.