diff --git a/.github/actions/unit_test/action.yml b/.github/actions/unit_test/action.yml
index 00cb28386..ebcad6c80 100644
--- a/.github/actions/unit_test/action.yml
+++ b/.github/actions/unit_test/action.yml
@@ -46,6 +46,7 @@ runs:
         export TF2ONNX_SKIP_TFLITE_TESTS=${{ inputs.skip_tflite }}
         export TF2ONNX_SKIP_TFJS_TESTS=True
         export TF2ONNX_SKIP_TF_TESTS=False
+        export TF_USE_LEGACY_KERAS=True
         python -m pytest --cov=tf2onnx --cov-report=term --disable-pytest-warnings -r s tests --cov-append --junitxml=junit/test-results.xml
         ls
 
@@ -58,5 +59,6 @@ runs:
         set TF2ONNX_SKIP_TFLITE_TESTS=${{ inputs.skip_tflite }}
         set TF2ONNX_SKIP_TFJS_TESTS=True
         set TF2ONNX_SKIP_TF_TESTS=False
+        set TF_USE_LEGACY_KERAS=True
         python -m pytest --cov=tf2onnx --cov-report=term --disable-pytest-warnings -r s tests --cov-append --junitxml=junit/test-results.xml
         ls
diff --git a/tests/utils/setup_test_env.sh b/tests/utils/setup_test_env.sh
index a14828d05..d74df3447 100755
--- a/tests/utils/setup_test_env.sh
+++ b/tests/utils/setup_test_env.sh
@@ -28,6 +28,8 @@ pip install tensorflow==$TF_VERSION
 pip uninstall -y protobuf
 pip install "protobuf~=3.20"
 
+pip install tf_keras
+
 python setup.py install
 
 echo "----- List all of depdencies:"