diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index 84aa5ae4a8af..c355e6e99e2a 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -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 =