Is the implementation wrong?

[ranery]

sorted_indices = torch.argsort(topk_indices)

should be

sorted_indices = torch.sort(topk_indices)

Right?

[sramshetty]

Nice catch, sorry about that!

I do in fact need the argsort() output for the gather here, but like you mentioned the gather for x and y should be the sorted indices themselves.

Let me know if this makes sense to you, and thanks for pointing it out.

[ranery]

That makes sense now!

I also used:

sorted_indices, index = torch.sort(topk_indices, dim=1)

so that the sorted_indices could be used for x gather and the index could be used for output gather.