zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-15 07:28:09 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Eval before cleanup so model file is unlocked (#1702)

Browse Source
This commit is contained in:
Cheng
2024-12-15 14:41:49 +09:00
committed by GitHub
parent f9640e049d
commit af5a614aad
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/tests/test_nn.py
View File

@@ -95,6 +95,9 @@ class TestBase(mlx_tests.MLXTestCase):
m.save_weights(npz_file)
m_load = make_model()
m_load.load_weights(npz_file)
# Eval before cleanup so model file is unlocked.
mx.eval(m_load.state)
tdir.cleanup()
eq_tree = tree_map(mx.array_equal, m.parameters(), m_load.parameters())
@@ -110,6 +113,9 @@ class TestBase(mlx_tests.MLXTestCase):
m.save_weights(safetensors_file)
m_load = make_model()
m_load.load_weights(safetensors_file)
# Eval before cleanup so model file is unlocked.
mx.eval(m_load.state)
tdir.cleanup()
eq_tree = tree_map(mx.array_equal, m.parameters(), m_load.parameters())