Skip to content

Commit

Permalink
Use temporary folder for testing model saving in file editor (#20439)
Browse files Browse the repository at this point in the history
  • Loading branch information
sampathweb authored Nov 1, 2024
1 parent 57caffe commit f9ea1a0
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions keras/src/saving/file_editor_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,14 +98,16 @@ def test_scalar_weight(self):
model.add(keras.layers.Dense(1, activation="sigmoid", name="my_dense"))
model.compile(optimizer="adam", loss="mse", metrics=["mae"])
model.fit(np.array([[1]]), np.array([[1]]), verbose=0)
model.save("model.keras")
model.save_weights("model.weights.h5")
model_fpath = os.path.join(self.get_temp_dir(), "model.keras")
weights_fpath = os.path.join(self.get_temp_dir(), "model.weights.h5")
model.save(model_fpath)
model.save_weights(weights_fpath)

model_editor = KerasFileEditor("model.keras")
model_editor = KerasFileEditor(model_fpath)
self.assertEqual(
len(keras.src.tree.flatten(model_editor.weights_dict)), 8
)
model_weights_editor = KerasFileEditor("model.weights.h5")
model_weights_editor = KerasFileEditor(weights_fpath)
self.assertEqual(
len(keras.src.tree.flatten(model_weights_editor.weights_dict)), 8
)

0 comments on commit f9ea1a0

Please sign in to comment.