Does faiss index on gpu error out for anyone else?

[ssharpe42]

When trying to use the dataset indexing on a GPU, it fails below.

I have faiss-gpu==1.7.2 installed and datasets==2.1.0.

File /opt/conda/envs/labeling/lib/python3.9/site-packages/datasets/search.py:320, in FaissIndex._faiss_index_to_device(index, device)
    317 if isinstance(device, int):
    318     # Positive integers are directly mapped to GPU ids
    319     if device > -1:
--> 320         faiss_res = faiss.StandardGpuResources()
    321         index = faiss.index_cpu_to_gpu(faiss_res, device, index)
    322     # And negative integers mean using all GPUs
    323     else:

NameError: name 'faiss' is not defined

If I don't specify a device it works on the cpu. Could this relate to the fact that faiss isn't imported inside _faiss_index_to_device? I see it imported in other functions but not here

datasets/src/datasets/search.py

Line 305 in b564af7

def _faiss_index_to_device(index: "faiss.Index", device: Optional[Union[int, List[int]]] = None) -> "faiss.Index":

If I add the import, the error goes away.

[alvarobartt]

Hi @ssharpe42 this was just solved in 🤗 datasets==2.2.0 available at 2.2.0.

More information at #4287 and PR at #4288

[alvarobartt]

@mariosasko could we already mark this discussion as answered? Thanks 🤗

[ssharpe42]

Great thanks!