At the momnet I found faiss 283 package to be useful.
Here’s an example 355 of using faiss with pytorch.
It’s not perfect (takes a lot of GPU memory, crashes every once in a while) but it does the job for me.