diff --git a/WINDOWS_INSTALL.bat b/WINDOWS_INSTALL.bat index 81a1f02..d414e87 100644 --- a/WINDOWS_INSTALL.bat +++ b/WINDOWS_INSTALL.bat @@ -310,7 +310,7 @@ for /f %%i in ('%PYTHON_CONSOLE_PATH% -c "from abogen.is_nvidia import check; pr echo. echo Checking CUDA availability... if /I "%IS_NVIDIA%"=="true" ( - for /f %%i in ('%PYTHON_CONSOLE_PATH% -c "from torch.cuda import is_available; print(is_available())"') do set cuda_available=%%i + for /f %%i in ('%PYTHON_CONSOLE_PATH% %PROJECTFOLDER%\check_cuda.py') do set cuda_available=%%i if "%cuda_available%"=="False" ( echo "Installing PyTorch with CUDA (12.8) support..." diff --git a/abogen/check_cuda.py b/abogen/check_cuda.py new file mode 100644 index 0000000..e674015 --- /dev/null +++ b/abogen/check_cuda.py @@ -0,0 +1,30 @@ +import sys +import os +import platform +import ctypes +import importlib.util + +def check_cuda_with_fix(): + """ + Check if CUDA is available, with a fix for PyTorch DLL loading issue + ([WinError 1114]) on Windows. + """ + # Fix PyTorch DLL loading issue ([WinError 1114]) on Windows + try: + if platform.system() == "Windows": + spec = importlib.util.find_spec("torch") + if spec and spec.origin: + dll_path = os.path.join(os.path.dirname(spec.origin), "lib", "c10.dll") + if os.path.exists(dll_path): + ctypes.CDLL(os.path.normpath(dll_path)) + except Exception: + pass + + try: + from torch.cuda import is_available + print(is_available()) + except ImportError: + print("False") + +if __name__ == "__main__": + check_cuda_with_fix()