Skip to content

Commit

Permalink
Change the way to call legacy functions.
Browse files Browse the repository at this point in the history
Signed-off-by: Jay Zhang <jiz@microsoft.com>
  • Loading branch information
fatcat-z committed Nov 25, 2024
1 parent c68331d commit 0d6c98e
Showing 1 changed file with 9 additions and 8 deletions.
17 changes: 9 additions & 8 deletions tests/keras2onnx_unit_tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,6 @@
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():
np.random.seed(42)
Expand All @@ -31,10 +27,15 @@ def runner():
def runner_func(*args, **kwargs):
return run_onnx_runtime(*args, model_files, **kwargs)

# Ensure Keras layer naming is reset for each function
K.reset_uids()
# Reset the TensorFlow session to avoid resource leaking between tests
K.clear_session()
if is_keras_3():
import tf_keras
tf_keras.backend.reset_uids()
tf_keras.backend.clear_session()
else:
# Ensure Keras layer naming is reset for each function
K.reset_uids()
# Reset the TensorFlow session to avoid resource leaking between tests
K.clear_session()

# Provide wrapped run_onnx_runtime function
yield runner_func
Expand Down

0 comments on commit 0d6c98e

Please sign in to comment.