import faiss
import torch

my_index = faiss.read_index("knn.index")


query = torch.rand(3, 1024)

k = 5
distances, indices = my_index.search(query, k)

import pdb;pdb.set_trace()