From a229c8cef03ef71e51fb3c8e3f86f4cefcee188a Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Tue, 4 Feb 2025 13:20:57 -0800
Subject: [PATCH] don't duplicate malloc with custom kernel init (#1830)

---
 mlx/backend/metal/custom_kernel.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp
index 66f75a65f..e775f6798 100644
--- a/mlx/backend/metal/custom_kernel.cpp
+++ b/mlx/backend/metal/custom_kernel.cpp
@@ -15,10 +15,11 @@ void CustomKernel::eval_gpu(
   std::vector<array> copies;
 
   for (auto& out : outputs) {
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
     if (init_value_) {
       copies.emplace_back(init_value_.value(), out.dtype());
       fill_gpu(copies.back(), out, s);
+    } else {
+      out.set_data(allocator::malloc_or_wait(out.nbytes()));
     }
   }