diff --git a/tests/test_loraplus_helper.py b/tests/test_loraplus_helper.py index 92b727342a7..ad8f068a777 100644 --- a/tests/test_loraplus_helper.py +++ b/tests/test_loraplus_helper.py @@ -1,9 +1,15 @@ -import bitsandbytes as bnb import torch from torch import nn +from peft.import_utils import is_bnb_available from peft.optimizers import create_loraplus_optimizer +from .testing_utils import require_bitsandbytes + + +if is_bnb_available(): + import bitsandbytes as bnb + class SimpleNet(nn.Module): def __init__(self, bias=True): @@ -21,6 +27,7 @@ def forward(self, X): return X +@require_bitsandbytes def test_lora_plus_helper_sucess(): model = SimpleNet() optimizer_cls = bnb.optim.Adam8bit @@ -46,6 +53,7 @@ def test_lora_plus_helper_sucess(): assert optim.param_groups[2]["lr"] == optim.param_groups[3]["lr"] == (optim_config["lr"] * loraplus_lr_ratio) +@require_bitsandbytes def test_lora_plus_optimizer_sucess(): """ Test if the optimizer is correctly created and step function runs without any exception