From 14efd9c35a72fe96f6edd963cedacdc2f5ef5493 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Mon, 18 Aug 2025 22:46:21 -0700
Subject: [PATCH] Fix compilation

---
 mlx/backend/cuda/custom_kernel.cpp  | 2 +-
 mlx/backend/metal/custom_kernel.cpp | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlx/backend/cuda/custom_kernel.cpp b/mlx/backend/cuda/custom_kernel.cpp
index c66fc56d4..d475dbc9f 100644
--- a/mlx/backend/cuda/custom_kernel.cpp
+++ b/mlx/backend/cuda/custom_kernel.cpp
@@ -334,7 +334,7 @@ void CustomKernel::eval_gpu(
       args.append_ndim(in.strides());
     }
     if (shape_info.ndim) {
-      args.append(in.ndim());
+      args.append<int32_t>(in.ndim());
     }
   }
   for (auto& out : outputs) {
diff --git a/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp
index e40d853fc..41e399ce3 100644
--- a/mlx/backend/metal/custom_kernel.cpp
+++ b/mlx/backend/metal/custom_kernel.cpp
@@ -172,7 +172,7 @@ std::string write_template(
   return template_def.str();
 }
 
-MetalKernelFunction metal_kernel(
+CustomKernelFunction metal_kernel(
     const std::string& name,
     const std::vector<std::string>& input_names,
     const std::vector<std::string>& output_names,