diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp
index d21c66ade..0a69dd261 100644
--- a/mlx/backend/metal/allocator.cpp
+++ b/mlx/backend/metal/allocator.cpp
@@ -272,9 +272,13 @@ Buffer MetalAllocator::malloc(size_t size) {
     if (!buf) {
       buf = device_->newBuffer(size, resource_options);
     }
+    if (!buf) {
+      return Buffer{nullptr};
+    }
     lk.lock();
-    if (buf) {
-      num_resources_++;
+    num_resources_++;
+    if (!buf->heap()) {
+      residency_set_.insert(buf);
     }
   }
 
@@ -288,10 +292,6 @@ Buffer MetalAllocator::malloc(size_t size) {
         get_cache_memory() - max_pool_size_);
   }
 
-  if (!buf->heap()) {
-    residency_set_.insert(buf);
-  }
-
   return Buffer{static_cast<void*>(buf)};
 }