From a002797d527c0846dfc9de1dc4c13009758bd02f Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Thu, 21 Dec 2023 17:59:15 -0800
Subject: [PATCH] A temporary fix (#254)

---
 mlx/backend/metal/device.cpp | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp
index 68662763a..eda54ea89 100644
--- a/mlx/backend/metal/device.cpp
+++ b/mlx/backend/metal/device.cpp
@@ -108,6 +108,16 @@ MTL::Library* load_library(
   }
 }
 
+struct PoolHolder {
+  PoolHolder() {
+    p = NS::AutoreleasePool::alloc()->init();
+  }
+  ~PoolHolder() {
+    p->release();
+  }
+  NS::AutoreleasePool* p;
+};
+
 } // namespace
 
 Device::Device()
@@ -125,6 +135,12 @@ Device::~Device() {
   for (auto& l : library_map_) {
     l.second->release();
   }
+  for (auto& b : buffer_map_) {
+    b.second.second->release();
+  }
+  for (auto& e : encoder_map_) {
+    e.second->release();
+  }
   device_->release();
   pool_->release();
 }
@@ -282,9 +298,8 @@ Device& device(mlx::core::Device) {
 }
 
 NS::AutoreleasePool*& thread_autorelease_pool() {
-  static thread_local NS::AutoreleasePool* p =
-      NS::AutoreleasePool::alloc()->init();
-  return p;
+  static thread_local PoolHolder pool{};
+  return pool.p;
 }
 
 void new_stream(Stream stream) {