torch: enable autoAddCudaCompatRunpath on Jetsons (#378960)

2025-02-02 23:44:18 +01:00
parent 2c704b65f9 94ad953c22
commit 56291f2eb3
1 changed files with 3 additions and 0 deletions
@@ -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 =