summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorTim Dettmers <tim.dettmers@gmail.com>2022-08-07 09:49:36 -0700
committerTim Dettmers <tim.dettmers@gmail.com>2022-08-07 09:49:36 -0700
commit26efb154c8d77b4ede2cfc0dbd2381dd385f33e7 (patch)
tree5f028f5405c02154f8420dfc794f16847a170126
parenta4532c59f7e6c2413838e740124aacb6e801f13b (diff)
Fixed bug where python -m bitsandbytes was failing.
-rw-r--r--bitsandbytes/__main__.py5
1 files changed, 3 insertions, 2 deletions
diff --git a/bitsandbytes/__main__.py b/bitsandbytes/__main__.py
index 7f3d24c..5f11875 100644
--- a/bitsandbytes/__main__.py
+++ b/bitsandbytes/__main__.py
@@ -30,7 +30,7 @@ print()
from . import COMPILED_WITH_CUDA, PACKAGE_GITHUB_URL
-from .cuda_setup.main import get_compute_capabilities
+from .cuda_setup.main import get_compute_capabilities, get_cuda_lib_handle
from .cuda_setup.env_vars import to_be_ignored
from .utils import print_stderr
@@ -47,7 +47,8 @@ print(
print_header("OTHER")
print(f"{COMPILED_WITH_CUDA = }")
-print(f"COMPUTE_CAPABILITIES_PER_GPU = {get_compute_capabilities()}")
+cuda = get_cuda_lib_handle()
+print(f"COMPUTE_CAPABILITIES_PER_GPU = {get_compute_capabilities(cuda)}")
print_header("")
print_header("DEBUG INFO END")
print_header("")