Is there a proper way to run k-NN function on GPU with a user defined distance function in Python?

Question

I want to run k-Nearest Neighbor (k-NN) with a new user defined distance function in Python. Could you please advice a proper way for it?

I could not find a solution how to run k-Nearest Neighbor (k-NN) with a new user defined distance function in Python.

I also installed Ubuntu (wsl2), because my windows did not see my GPU device.

I even tried to calculate the distance matrix with cupy then wanted to transfer it scikitlearn knnClassifier. However, calculating the image distance matrix takes too long even on GPU using nested loops.

...
import cupy
from sklearn.neighbors import NearestNeighbors as skNearestNeighbors
...

Our x_train and x_test are cupy.ndarray. x_train.shape is 16384.

We are trying to calculate dist matrix on GPU, however by using the code below takes approximately 15 mins. for each iteration.

num_samples = len(x_train)
distance_matrix = cupy.empty((len(x_train), len(x_test)), dtype=cupy.float32)


for i in range(len(x_train)):
    for j in range(len(x_test)):
        distance_matrix[i, j] = user_defined_distFunction(x_train[i, ], x_test[j, ])

After calculating this distance_matrix we are planning to give it as the input for scikitlearn's knnClassifier.

So, is there a proper way (any library etc.) to run k-NN function on GPU with a user defined distance function in Python?