Pure PyTorch implementation of KNN with both CPU and GPU versions. Probably not as fast as CUDA implementations, but written in pure PyTorch so it'll have less issues (no need to specify Python or CUDA paths).
Implementation is in progress (not working).