[`bnb`] Add simple check for bnb import (#24995)

add simple check for bnb

[`bnb`] Add simple check for bnb import (#24995)
add simple check for bnb
d3ce048c · Younes Belkada · GitHub · f1a1eb4a · d3ce048c
Unverified Commit d3ce048c authored Jul 21, 2023 by Younes Belkada Committed by GitHub Jul 21, 2023
Show whitespace changes
Inline Side-by-side

Showing with 8 additions and 1 deletion

src/transformers/utils/import_utils.py src/transformers/utils/import_utils.py +8 -1

No files found.
--- a/src/transformers/utils/import_utils.py
+++ b/src/transformers/utils/import_utils.py
@@ -504,7 +504,14 @@ def is_ipex_available():


 def is_bitsandbytes_available():
-    return _bitsandbytes_available
+    if not is_torch_available():
+        return False
+
+    # bitsandbytes throws an error if cuda is not available
+    # let's avoid that by adding a simple check
+    import torch
+
+    return _bitsandbytes_available and torch.cuda.is_available()


 def is_torchdistx_available():