Update keras.backend for keras 3.

Signed-off-by: Jay Zhang <[email protected]>
onnx · Nov 25, 2024 · 77e5c39 · 77e5c39
1 parent 445a89e
commit 77e5c39
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tests/keras2onnx_unit_tests/conftest.py b/tests/keras2onnx_unit_tests/conftest.py
@@ -12,6 +12,12 @@
 
 K = keras.backend
 
+def is_keras_3():
+    return tf.__version__.startswith("2.18") or tf.__version__.startswith("2.17") or tf.__version__.startswith("2.16")
+
+if is_keras_3():
+    import tf_keras
+    K = tf_keras.backend
 
 @pytest.fixture(scope='function')
 def runner():