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,