torch: enable autoAddCudaCompatRunpath on Jetsons

This commit is contained in:
Conroy Cheers
2025-02-02 23:32:23 +01:00
committed by Pavol Rusnak
parent 508be1c431
commit 94ad953c22
@@ -150,6 +150,8 @@ let
supportedCudaCapabilities = lists.intersectLists cudaFlags.cudaCapabilities supportedTorchCudaCapabilities;
unsupportedCudaCapabilities = lists.subtractLists supportedCudaCapabilities cudaFlags.cudaCapabilities;
isCudaJetson = cudaSupport && cudaPackages.cudaFlags.isJetsonBuild;
# Use trivial.warnIf to print a warning if any unsupported GPU targets are specified.
gpuArchWarner =
supported: unsupported:
@@ -456,6 +458,7 @@ buildPythonPackage rec {
cuda_nvcc
]
)
++ lib.optionals isCudaJetson [ cudaPackages.autoAddCudaCompatRunpath ]
++ lib.optionals rocmSupport [ rocmtoolkit_joined ];
buildInputs =