zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-30 06:31:21 +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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 0 deletions
Show Stats 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())