fix import (#10103)

226973a9 · Patrick von Platen · GitHub · 4cda2d73 · 226973a9
Unverified Commit 226973a9 authored Feb 09, 2021 by Patrick von Platen Committed by GitHub Feb 09, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

src/transformers/file_utils.py src/transformers/file_utils.py +5 -1

No files found.
--- a/src/transformers/file_utils.py
+++ b/src/transformers/file_utils.py
@@ -144,7 +144,11 @@ try:
    _faiss_version = importlib_metadata.version("faiss")
    logger.debug(f"Successfully imported faiss version {_faiss_version}")
 except importlib_metadata.PackageNotFoundError:
-    _faiss_available = False
+    try:
+        _faiss_version = importlib_metadata.version("faiss-cpu")
+        logger.debug(f"Successfully imported faiss version {_faiss_version}")
+    except importlib_metadata.PackageNotFoundError:
+        _faiss_available = False


 _scatter_available = importlib.util.find_spec("torch_scatter") is not None