mirror of
https://github.com/ml-explore/mlx.git
synced 2025-11-05 11:28:12 +08:00
move memory APIs into top level mlx.core (#1982)
This commit is contained in:
@@ -635,7 +635,7 @@ class TestVmap(mlx_tests.MLXTestCase):
|
||||
|
||||
def test_leaks(self):
|
||||
if mx.metal.is_available():
|
||||
mem_pre = mx.metal.get_active_memory()
|
||||
mem_pre = mx.get_active_memory()
|
||||
else:
|
||||
mem_pre = 0
|
||||
|
||||
@@ -653,7 +653,7 @@ class TestVmap(mlx_tests.MLXTestCase):
|
||||
gc.collect()
|
||||
|
||||
if mx.metal.is_available():
|
||||
mem_post = mx.metal.get_active_memory()
|
||||
mem_post = mx.get_active_memory()
|
||||
else:
|
||||
mem_post = 0
|
||||
|
||||
|
||||
Reference in New Issue
Block a user