diff --git a/bitsandbytes/cuda_setup/main.py b/bitsandbytes/cuda_setup/main.py index e7901d82e..562baa44b 100644 --- a/bitsandbytes/cuda_setup/main.py +++ b/bitsandbytes/cuda_setup/main.py @@ -111,7 +111,19 @@ def run_cuda_setup(self): package_dir = Path(__file__).parent.parent binary_path = package_dir / binary_name - print('bin', binary_path) + if not binary_path.exists(): + # If binary file does not exist, check LD_LIBRARY_PATH + library_paths = os.getenv("LD_LIBRARY_PATH", "").split(":") + print("library_paths ", library_paths) + for path in library_paths: + potential_binary_path = Path(path) / binary_name + print("potential_binary_path", potential_binary_path) + if potential_binary_path.exists(): + binary_path = potential_binary_path + print("binary_path", binary_path) + break + + print("bin", binary_path) try: if not binary_path.exists():