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

ml-explore · Dec 15, 2024 · af5a614 · af5a614
1 parent f9640e0
commit af5a614
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py
@@ -95,6 +95,9 @@ def make_model():
         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 @@ def make_model():
         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())