From b2d2b37888fe12a5901266369236e207ed3faa30 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Sun, 30 Mar 2025 16:27:26 -0700
Subject: [PATCH] fix residency set clearing (#2027)

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

diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp
index 9911acef1..2c81482a5 100644
--- a/mlx/backend/metal/allocator.cpp
+++ b/mlx/backend/metal/allocator.cpp
@@ -48,6 +48,9 @@ int BufferCache::clear() {
   int n_release = 0;
   for (auto& [size, holder] : buffer_pool_) {
     if (holder->buf) {
+      if (holder->buf->heap()) {
+        residency_set_.erase(holder->buf);
+      }
       holder->buf->release();
       n_release++;
     }