From d3ce048c2029c829e23c7d82f9abd4919610db87 Mon Sep 17 00:00:00 2001 From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Date: Fri, 21 Jul 2023 17:50:52 +0200 Subject: [PATCH] [`bnb`] Add simple check for bnb import (#24995) add simple check for bnb --- src/transformers/utils/import_utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index ffefaed9e62..c0a8c80f0b0 100644 --- 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():