Skip to content

Commit ca2275d

Browse files
committed
Fix build issue. CUDA may be installed in $CUDA_HOME/lib on macOS.
Closes gh-16955. Closes pytorch/vision#977 On Linux both `lib64` and `lib` may be present (symlinked). The reports seem to all be about macOS, but it seems like this is also possibly more robust on Linux and can't hurt. So not treating platforms differently. Note that Eigen has a similar check in its CMake: ``` if(CUDA_64_BIT_DEVICE_CODE AND (EXISTS "${CUDA_TOOLKIT_ROOT_DIR}/lib64")) link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") else() link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib") endif() ``` There may be other issues for building from source on macOS, can't test.
1 parent ef63561 commit ca2275d

File tree

1 file changed

+11
-1
lines changed

1 file changed

+11
-1
lines changed

torch/utils/cpp_extension.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -548,7 +548,17 @@ def library_paths(cuda=False):
548548
paths.append(lib_path)
549549

550550
if cuda:
551-
lib_dir = 'lib/x64' if IS_WINDOWS else 'lib64'
551+
if IS_WINDOWS:
552+
lib_dir = 'lib/x64'
553+
else:
554+
lib_dir = 'lib64'
555+
if (not os.path.exists(_join_cuda_home(lib_dir)) and
556+
os.path.exists(_join_cuda_home('lib'))):
557+
# 64-bit CUDA may be installed in 'lib' (see e.g. gh-16955)
558+
# Note that it's also possible both don't exist (see
559+
# _find_cuda_home) - in that case we stay with 'lib64'.
560+
lib_dir = 'lib'
561+
552562
paths.append(_join_cuda_home(lib_dir))
553563
if CUDNN_HOME is not None:
554564
paths.append(os.path.join(CUDNN_HOME, lib_dir))

0 commit comments

Comments
 (0)