From 46356c61457d662c861db5ecb056d2dc14e734a2 Mon Sep 17 00:00:00 2001
From: Yuxuan Shui <yshuiv7@gmail.com>
Date: Thu, 4 Jul 2024 17:52:35 +0100
Subject: [PATCH] ollama: fix cuda support

Fixes: 82018339bd1a587e63f5a12c8517252f334220e0
---
 pkgs/by-name/ol/ollama/package.nix | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pkgs/by-name/ol/ollama/package.nix b/pkgs/by-name/ol/ollama/package.nix
index 93c1198c5b0e..c8f7f7beccc4 100644
--- a/pkgs/by-name/ol/ollama/package.nix
+++ b/pkgs/by-name/ol/ollama/package.nix
@@ -142,6 +142,8 @@ in
 goBuild ((lib.optionalAttrs enableRocm {
   ROCM_PATH = rocmPath;
   CLBlast_DIR = "${clblast}/lib/cmake/CLBlast";
+}) // (lib.optionalAttrs enableCuda {
+  CUDA_LIB_DIR = "${cudaToolkit}/lib";
 }) // {
   inherit pname version src vendorHash;
 
@@ -161,6 +163,7 @@ goBuild ((lib.optionalAttrs enableRocm {
     (rocmLibs ++ [ libdrm ])
   ++ lib.optionals enableCuda [
     cudaPackages.cuda_cudart
+    cudaPackages.cuda_cccl
     cudaPackages.libcublas
   ] ++ lib.optionals stdenv.isDarwin
     metalFrameworks;