Fix exception caused by TensorFlow 2.5 KerasTensor API incompatible change
@@ -345,11 +345,14 @@ def keras_session(model, input_tensor_names, output_tensor_names, **kwargs):
if not isinstance(model, tf.keras.Model):
model = tf.keras.models.load_model(model)
kwargs = dict(zip(model.input_names, model.inputs))
- if tf.version.VERSION > '2.2.0':
+ if tf.version.VERSION > '2.2.0' and tf.version.VERSION < '2.5.0':
from tensorflow.python.keras.engine import keras_tensor
if keras_tensor.keras_tensors_enabled():
for name, tensor in kwargs.items():
kwargs[name] = tensor.type_spec
+ elif tf.version.VERSION >= '2.5.0':
+ for name, tensor in kwargs.items():
+ kwargs[name] = tensor.type_spec
full_model = tf.function(lambda **kwargs: model(kwargs.values()))
concrete_function = full_model.get_concrete_function(**kwargs)
frozen_model = convert_variables_to_constants_v2(concrete_function)
