Skip to content

Commit 9a0a226

Browse files
committed
Fix bitsandbytes imports in unit tests on MacOS.
1 parent 477d87e commit 9a0a226

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_custom_invoke_linear_nf4.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
import pytest
22
import torch
33

4-
from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_invoke_linear_nf4 import (
5-
CustomInvokeLinearNF4,
6-
)
74
from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
85
wrap_custom_layer,
96
)
10-
from invokeai.backend.quantization.bnb_nf4 import InvokeLinearNF4
7+
8+
if not torch.cuda.is_available():
9+
pytest.skip("CUDA is not available", allow_module_level=True)
10+
else:
11+
from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_invoke_linear_nf4 import (
12+
CustomInvokeLinearNF4,
13+
)
14+
from invokeai.backend.quantization.bnb_nf4 import InvokeLinearNF4
1115

1216

1317
def build_linear_nf4_layer(orig_layer: torch.nn.Linear | None = None):

0 commit comments

Comments
 (0)