Fix digit knn function