From f5cd03c64dae58b997cfb241720c6f5b280ba9ca Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Tue, 5 Nov 2024 17:04:46 -0800
Subject: [PATCH] fix stream

---
 llms/mlx_lm/utils.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llms/mlx_lm/utils.py b/llms/mlx_lm/utils.py
index 9f852ae3..853e1c95 100644
--- a/llms/mlx_lm/utils.py
+++ b/llms/mlx_lm/utils.py
@@ -35,8 +35,7 @@ MODEL_REMAPPING = {
 MAX_FILE_SIZE_GB = 5
 
 # A stream on the default device just for generation
-# generation_stream = mx.new_stream(mx.default_device())
-generation_stream = mx.default_stream(mx.default_device())
+generation_stream = mx.new_stream(mx.default_device())
 
 
 class ModelNotFoundError(Exception):