Fix GPU memory contention between jax and enabling tf_eager
Fix GPU memory contention between jax and enabling tf_eager

PiperOrigin-RevId: 262156649

tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor