Avoid use of deprecated xla_bridge.get_backend().platform

Co-authored-by: Jake Vanderplas <[email protected]>
huggingface · Feb 6, 2025 · 52d10f6 · 52d10f6
1 parent e3d7298
commit 52d10f6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/commands/env.py b/src/transformers/commands/env.py
@@ -116,7 +116,7 @@ def run(self):
             flax_version = flax.__version__
             jax_version = jax.__version__
             jaxlib_version = jaxlib.__version__
-            jax_backend = jax.lib.xla_bridge.get_backend().platform
+            jax_backend = jax.default_backend()
 
         info = {
             "`transformers` version": version,