From 855b9881f07ae0227c9a3773a69b9bd2ec7e4602 Mon Sep 17 00:00:00 2001 From: zehao-intel Date: Mon, 27 May 2024 15:47:32 +0800 Subject: [PATCH] Fix INC TF 2x Estimator Issue on TF2.16 (#1819) Signed-off-by: zehao-intel --- neural_compressor/model/tensorflow_model.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/neural_compressor/model/tensorflow_model.py b/neural_compressor/model/tensorflow_model.py index e4809863a55..0caadc53dbd 100644 --- a/neural_compressor/model/tensorflow_model.py +++ b/neural_compressor/model/tensorflow_model.py @@ -81,7 +81,9 @@ def get_model_type(model): return "graph" elif isinstance(model, tf.compat.v1.GraphDef): return "graph_def" - elif isinstance(model, tf.compat.v1.estimator.Estimator): + elif not version1_gte_version2(tf.version.VERSION, "2.16.1") and isinstance( + model, tf.compat.v1.estimator.Estimator + ): return "estimator" elif isinstance(model, str): model = os.path.abspath(os.path.expanduser(model))