I would like to create a knn graph of this tensor a using torch such that it returns me k indices and distances, for each row of this a tensor. That is to say distances.shape is [10,k] and indices.shape is [10,k]

Basically I look to do something like this(section 1.6.1.1. Finding the Nearest Neighbors) using torch. It is not very evident to me how to do this?

distances = dist.values[:,0:k]
indices = dist.indices[:,0:k]

This way I have what I wanted, but perhaps there are more efficient ways to achieve this?
PS: This method is highly memory inefficient is kd-tree calculation possible using torch?