From 6cf48872b768d25aa29c52fa6b9a36f19c7025f0 Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Sun, 6 Apr 2025 12:05:16 +0900
Subject: [PATCH] wait_for_one should wait for task to finish (#2047)

---
 mlx/scheduler.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlx/scheduler.h b/mlx/scheduler.h
index bf34b38c0..b2c6b842b 100644
--- a/mlx/scheduler.h
+++ b/mlx/scheduler.h
@@ -129,7 +129,7 @@ class Scheduler {
     int n_tasks_old = n_active_tasks();
     if (n_tasks_old > 1) {
       completion_cv.wait(lk, [this, n_tasks_old] {
-        return this->n_active_tasks() != n_tasks_old;
+        return this->n_active_tasks() < n_tasks_old;
       });
     }
   }