From 0e6724e0535d2330c3a30261c7076bf6eabd0d54 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 28 Apr 2026 13:07:49 +0000
Subject: [PATCH] python3Packages.executorch: fix cuda build

---
 .../python-modules/executorch/default.nix     | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/pkgs/development/python-modules/executorch/default.nix b/pkgs/development/python-modules/executorch/default.nix
index c8e393dc4166..02604cfe8f10 100644
--- a/pkgs/development/python-modules/executorch/default.nix
+++ b/pkgs/development/python-modules/executorch/default.nix
@@ -7,6 +7,8 @@
 
   # nativeBuildInputs
   gitMinimal,
+  # cuda-only:
+  autoPatchelfHook,
 
   # build-system
   certifi,
@@ -46,8 +48,11 @@
   transformers,
   writableTmpDirAsHomeHook,
   yaspin,
+
+  cudaSupport ? torch.cudaSupport,
+  cudaPackages,
 }:
-buildPythonPackage (finalAttrs: {
+buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
   pname = "executorch";
   version = "1.2.0";
   pyproject = true;
@@ -107,6 +112,9 @@ buildPythonPackage (finalAttrs: {
     #  Some binaries contain forbidden references to /build/. Check the error above!
     CMAKE_ARGS = lib.concatStringsSep " " [
       (lib.cmakeBool "CMAKE_SKIP_BUILD_RPATH" true)
+
+      # For some cmake-tier reason, cmakeBool does not work here
+      (lib.cmakeFeature "EXECUTORCH_BUILD_CUDA" (if cudaSupport then "ON" else "OFF"))
     ];
   };
 
@@ -122,6 +130,15 @@ buildPythonPackage (finalAttrs: {
 
   nativeBuildInputs = [
     gitMinimal
+  ]
+  ++ lib.optionals cudaSupport [
+    autoPatchelfHook
+    cudaPackages.cuda_nvcc
+  ];
+
+  buildInputs = lib.optionals cudaSupport [
+    cudaPackages.cuda_cudart
+    cudaPackages.cuda_nvrtc
   ];
 
   pythonRemoveDeps = [