add lib suffix ".dylib" for Darwin

Co-authored-by: Aarni Koskela <[email protected]>
bitsandbytes-foundation · Feb 6, 2024 · dfc462c · dfc462c
1 parent f4238a4
commit dfc462c
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/bitsandbytes/cuda_setup/main.py b/bitsandbytes/cuda_setup/main.py
@@ -28,16 +28,15 @@
 
 from .env_vars import get_potentially_lib_path_containing_env_vars
 
-if platform.system() == 'Windows':  # Windows
+DYNAMIC_LIBRARY_SUFFIX = { "Darwin": ".dylib", "Windows": ".dll", "Linux": ".so"}.get(platform.system(), ".so")
+if platform.system() == "Windows":  # Windows
     CUDA_RUNTIME_LIBS = ["nvcuda.dll"]
-    DYNAMIC_LIBRARY_SUFFIX = ".dll"
 else:  # Linux or other
     # these are the most common libs names
     # libcudart.so is missing by default for a conda install with PyTorch 2.0 and instead
     # we have libcudart.so.11.0 which causes a lot of errors before
     # not sure if libcudart.so.12.0 exists in pytorch installs, but it does not hurt
     CUDA_RUNTIME_LIBS = ["libcudart.so", "libcudart.so.11.0", "libcudart.so.12.0", "libcudart.so.12.1", "libcudart.so.12.2"]
-    DYNAMIC_LIBRARY_SUFFIX = ".so"
 
 
 class CUDASetup: