frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

[CUDA] Change slim-wheel libraries load order (#145638)

Browse Source 
There is no libnvjitlink in  CUDA-11.x , so attempts to load it first will abort the execution and prevent the script from preloading nvrtc

Fixes issues reported in https://github.com/pytorch/pytorch/pull/145614#issuecomment-2613107072

Pull Request resolved: https://github.com/pytorch/pytorch/pull/145638
Approved by: https://github.com/atalman, https://github.com/kit1980, https://github.com/malfet

Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
This commit is contained in:
Wei Wang
2025-01-24 22:00:56 +00:00
committed by PyTorch MergeBot
parent 615bdd9c81
commit 2a70de7e92
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
torch/__init__.py
View File

@ -324,9 +324,10 @@ def _load_global_deps() -> None:
# libtorch_global_deps.so always depends in cudart, check if its installed via wheel
if "nvidia/cuda_runtime/lib/libcudart.so" not in _maps:
return
# If all abovementioned conditions are met, preload nvjitlink and nvrtc
_preload_cuda_deps("nvjitlink", "libnvJitLink.so.*[0-9]")
# If all above-mentioned conditions are met, preload nvrtc and nvjitlink
# Please note that order are important for CUDA-11.8 , as nvjitlink does not exist there
_preload_cuda_deps("cuda_nvrtc", "libnvrtc.so.*[0-9]")
_preload_cuda_deps("nvjitlink", "libnvJitLink.so.*[0-9]")
except Exception:
pass