One way, which I am able to arrive is:
dist = torch.sort(torch.cdist(a,a),dim=1)

Then assuming k=3
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?