torchKNN

KNN implement in Pytorch 1.0 including both cpu version and gpu version

dependency

tested with ubuntu16.04 + cuda9.0 + python3 + pytorch nightly

git clone https://github.com/foolyc/torchKNN.git
cd torchKNN
sudo sudo python3 setup.py install

ref [batch, dim, num_ref]

query [batch, dim, num_query]

inds [batch, k, num_query]

knn(ref, query, inds)

Name		Name	Last commit message	Last commit date
Latest commit History 2 Commits
torch_ops		torch_ops
.gitignore		.gitignore
LICENSE		LICENSE
README.md		README.md
setup.py		setup.py
test.py		test.py