resolve comments

Signed-off-by: Chenyaaang <chenyangli@google.com>

Author: Chenyaaang

tpu_inference/worker/tpu_worker_jax.py

@@ -165,8 +165,8 @@ def init_device(self):
             if device_indexes is not None and len(device_indexes) > 0:
                 # Enforcing the devices sequence to be consistent with the specified device indexes
                 self.devices = [jax.local_devices()[i] for i in device_indexes]
-                all_devices = jax.local_devices()
-                device_dict = {device.id: device for device in all_devices}
+                all_local_devices = jax.local_devices()
+                device_dict = {device.id: device for device in all_local_devices}
                 self.devices = []
                 for device_index in device_indexes:
                     device = device_dict[device_index]