mirror of
https://github.com/ml-explore/mlx.git
synced 2025-11-01 16:58:08 +08:00
move memory APIs into top level mlx.core (#1982)
This commit is contained in:
@@ -747,7 +747,7 @@ class TestAutograd(mlx_tests.MLXTestCase):
|
||||
]:
|
||||
if mx.metal.is_available():
|
||||
mx.synchronize(mx.default_stream(mx.default_device()))
|
||||
mem_pre = mx.metal.get_active_memory()
|
||||
mem_pre = mx.get_active_memory()
|
||||
else:
|
||||
mem_pre = 0
|
||||
|
||||
@@ -765,7 +765,7 @@ class TestAutograd(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