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

nWEIdia · malfet · pytorchbot · commit d544858876bf · 2025-01-24T22:36:57.000Z
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 #145614 (comment) Pull Request resolved: #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> (cherry picked from commit 2a70de7)
diff --git a/torch/__init__.py b/torch/__init__.py
@@ -328,9 +328,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
 
