mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[jax2tf] test: fix jax serialization version tests

Browse Source
This commit is contained in:
Jake VanderPlas 2024-06-04 11:03:06 -07:00
parent 2333d5c7c3
commit ca784a09a3
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/experimental/jax2tf/tests/call_tf_test.py
View File

@ -1154,7 +1154,7 @@ class RoundTripToTfTest(tf_test_util.JaxToTfTestCase):
if version != version_override:
self.addCleanup(partial(jax.config.update,
"jax_serialization_version",
version_override))
version))
jax.config.update("jax_serialization_version", version_override)
logging.info(
"Using JAX serialization version %s",