KevinMusgrave
commented
2 years ago When the faiss index is on multiple GPUs, then the inputs have to be on CPU (see https://github.com/facebookresearch/faiss/issues/1997). Inside the function, faiss moves the tensors back to the GPU.
When the index is on a single GPU, then the input tensor can also be on GPU.
Maybe I can add some code to move the tensor to CPU only if the index is on multiple GPUs.
pytorch_metric_learning.utils.inference
in function try_gpu(line258), weather gpu_index or cpu_index, both call function add_to_index_and_search, but this function only use cpu for faiss search
def add_to_index_and_search(index, query, reference, k): if reference is not None: index.add(reference.float().cpu()) return index.search(query.float().cpu(), k)