From 8058709d5a4970c3132c755a5f9fef41fa0ec931 Mon Sep 17 00:00:00 2001
From: younesbelkada <younesbelkada@gmail.com>
Date: Mon, 3 Apr 2023 16:27:30 +0000
Subject: [PATCH] fix failing CIs

---
 setup.py                 |  2 +-
 tests/test_common_gpu.py | 10 +++++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 2ece62b..e61396b 100644
--- a/setup.py
+++ b/setup.py
@@ -18,7 +18,7 @@ extras = {}
 extras["quality"] = ["black ~= 22.0", "ruff>=0.0.241"]
 extras["docs_specific"] = ["hf-doc-builder"]
 extras["dev"] = extras["quality"] + extras["docs_specific"]
-extras["test"] = extras["dev"] + ["pytest", "pytest-xdist", "parameterized"]
+extras["test"] = extras["dev"] + ["pytest", "pytest-xdist", "parameterized", "datasets"]
 
 setup(
     name="peft",
diff --git a/tests/test_common_gpu.py b/tests/test_common_gpu.py
index 1e99222..7b5a39e 100644
--- a/tests/test_common_gpu.py
+++ b/tests/test_common_gpu.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import gc
+import importlib
 import unittest
 
 import pytest
@@ -20,11 +21,18 @@ import torch
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer, WhisperForConditionalGeneration
 
 from peft import LoraConfig, PeftModel, get_peft_model
-from peft.tuners.lora import Linear8bitLt
 
 from .testing_utils import require_bitsandbytes, require_torch_gpu, require_torch_multi_gpu
 
 
+def is_bnb_available():
+    return importlib.util.find_spec("bitsandbytes") is not None
+
+
+if is_bnb_available():
+    from peft.tuners.lora import Linear8bitLt
+
+
 @require_torch_gpu
 class PeftGPUCommonTests(unittest.TestCase):
     r""" """