zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-12-15 01:42:31 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Clear cache every now and then (#1081)

Browse Source 
* clear cache every now and then

* don't need user arg anymore
This commit is contained in:
Awni Hannun
2024-11-01 14:15:32 -07:00
committed by GitHub
parent 8160e0c4e5
commit e510987870
2 changed files with 4 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
llms/mlx_lm/utils.py
View File

@@ -310,10 +310,14 @@ def generate_step(
y, logprobs = _step(y)
mx.async_eval(y, logprobs)
n = 0
while True:
next_y, next_logprobs = _step(y)
mx.async_eval(next_y, next_logprobs)
yield y.item(), logprobs
if n % 256 == 0:
mx.metal.clear_cache()
n += 1
y, logprobs = next_y, next_logprobs