From 65b22ef51b874b4669331879ff4323ba3ee10cb4 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Fri, 12 Jul 2024 17:48:25 +0800 Subject: [PATCH] fix: correct the params passed to keras.save_model (#4857) Signed-off-by: Frost Ming --- src/bentoml/_internal/frameworks/keras.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/bentoml/_internal/frameworks/keras.py b/src/bentoml/_internal/frameworks/keras.py index c91143a6a64..d5f7fdca6ce 100644 --- a/src/bentoml/_internal/frameworks/keras.py +++ b/src/bentoml/_internal/frameworks/keras.py @@ -251,6 +251,11 @@ def KerasSequentialModel() -> keras.models.Model: ) options = KerasOptions(include_optimizer=include_optimizer) + kwargs = {} + if tf_signatures is not None: + kwargs["signatures"] = tf_signatures + if tf_save_options is not None: + kwargs["options"] = tf_save_options with bentoml.models._create( # type: ignore name, @@ -264,12 +269,7 @@ def KerasSequentialModel() -> keras.models.Model: metadata=metadata, signatures=signatures, ) as bento_model: - model.save( - bento_model.path, - signatures=tf_signatures, - options=tf_save_options, - include_optimizer=include_optimizer, - ) + model.save(bento_model.path, include_optimizer=include_optimizer, **kwargs) return bento_model