fix: correct the params passed to keras.save_model (#4857)

Signed-off-by: Frost Ming <me@frostming.com>
bentoml · Jul 12, 2024 · 65b22ef · 65b22ef
1 parent bc9cca9
commit 65b22ef
Showing 1 changed file with 6 additions and 6 deletions.
diff --git 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