From cb2bba558975aca7720ba79c5f0db9081846624c Mon Sep 17 00:00:00 2001 From: Orbax Authors Date: Fri, 3 Jul 2026 01:31:18 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 942014089 --- export/orbax/export/modules/tensorflow_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export/orbax/export/modules/tensorflow_module.py b/export/orbax/export/modules/tensorflow_module.py index 99564565e..d1e59ad19 100644 --- a/export/orbax/export/modules/tensorflow_module.py +++ b/export/orbax/export/modules/tensorflow_module.py @@ -280,7 +280,7 @@ def jax_params_to_tf_variables( ) -> PyTree: """Converts `params` to tf.Variables in the same pytree structure.""" mesh = dtensor_utils.get_current_mesh() - default_cpu_device = tf.config.list_logical_devices('CPU')[0] + default_cpu_device = tf.config.list_logical_devices('CPU')[0].name if mesh is not None: if pspecs is None: raise ValueError(