From 94ad953c220dfd52d123b5885b324968312b3cb5 Mon Sep 17 00:00:00 2001
From: Conroy Cheers <conroy@dromeda.com.au>
Date: Fri, 25 Oct 2024 17:09:07 +1100
Subject: [PATCH] torch: enable autoAddCudaCompatRunpath on Jetsons

---
 pkgs/development/python-modules/torch/default.nix | 3 +++
 1 file changed, 3 insertions(+)

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 =