zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-05 03:18:12 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix memory count bug (#2717)

Browse Source
This commit is contained in:
Awni Hannun
2025-10-30 14:27:15 -07:00
committed by GitHub
parent 793a31eeb6
commit 68c5fa1c95
2 changed files with 15 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/tests/test_memory.py
View File

@@ -58,6 +58,20 @@ class TestMemory(mlx_tests.MLXTestCase):
with self.assertRaises(ValueError):
mx.set_wired_limit(max_size + 10)
def test_active_memory_count(self):
mx.synchronize()
mx.clear_cache()
init_mem = mx.get_active_memory()
a = mx.zeros((128, 128))
mx.eval(a)
mx.synchronize()
del a
a = mx.zeros((90, 128))
mx.eval(a)
mx.synchronize()
del a
self.assertEqual(init_mem, mx.get_active_memory())
if __name__ == "__main__":
mlx_tests.MLXTestRunner()