Reconstruction using PyTorch GPU tensors?

[ZwFink]

Using the Python interface, I am trying to use FAISS (1.7.3 installed via Pip) to find a point's k-nearest neighbors, then use the nearest neighbors in a downstream task without moving any data between the host and device. But it appears that index.reconstruct doesn't work on PyTorch tensors.

MWE:

import faiss
from faiss.contrib import torch_utils
import torch
import numpy as np

dim = 10
n_points = 1000
data = torch.randn(n_points, dim, device="cuda", dtype=torch.float32)

index = faiss.index_cpu_to_all_gpus(faiss.IndexFlatL2(dim))
index.add(data)  

query = torch.randn(100, dim, device="cuda", dtype=torch.float32)
D, I = index.search(query, k=5)

try:
    # I expect this to work and return a PyTorch tensor
    reconstructed = index.reconstruct_batch(I[0])
    print("Reconstructed:", reconstructed.shape)
except Exception as e:
    print("Failed reconstruction:", str(e))
    I_cpu = I.cpu().numpy()
    index.reconstruct_batch(I_cpu[0])
    print("Succeeded after moving to CPU")

It fails with this error:

Failed reconstruction: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
Succeeded after moving to CPU

[asadoughi]

Yes, it looks like reconstruct_batch is not wrapped in https://github.com/facebookresearch/faiss/blob/main/contrib/torch_utils.py. You can try using search_and_reconstruct as a workaround.