From 496278b63d90d7ecda2b0f9dac92dbe9656661e6 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <pavol@rusnak.io>
Date: Sat, 1 Mar 2025 17:52:03 +0100
Subject: [PATCH] katago: fix cuda and tensorrt builds

---
 pkgs/by-name/ka/katago/package.nix | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/pkgs/by-name/ka/katago/package.nix b/pkgs/by-name/ka/katago/package.nix
index 204d6777b0c4..e3ca8cd86781 100644
--- a/pkgs/by-name/ka/katago/package.nix
+++ b/pkgs/by-name/ka/katago/package.nix
@@ -34,8 +34,18 @@ assert lib.assertOneOf "backend" backend [
 let
   githash = "cd0ed6c0712088ddb901be68189ba7fa1439a9e7";
   fakegit = writeShellScriptBin "git" "echo ${githash}";
+  stdenv' =
+    if
+      builtins.elem backend [
+        "cuda"
+        "tensorrt"
+      ]
+    then
+      cudaPackages.backendStdenv
+    else
+      stdenv;
 in
-stdenv.mkDerivation rec {
+stdenv'.mkDerivation rec {
   pname = "katago";
   version = "1.15.3";
 
@@ -58,10 +68,12 @@ stdenv.mkDerivation rec {
     ]
     ++ lib.optionals (backend == "eigen") [ eigen ]
     ++ lib.optionals (backend == "cuda") [
+      cudaPackages.cuda_cudart
       cudaPackages.cudnn
       cudaPackages.cudatoolkit
     ]
     ++ lib.optionals (backend == "tensorrt") [
+      cudaPackages.cuda_cudart
       cudaPackages.cudatoolkit
       cudaPackages.tensorrt
     ]