This change enhances the TPU autodetection logic in notebook_launcher and remove call of device_count() before xmp.spawn for avoiding earlier devices initialising for passing check on torch_xla._XLAC._xla_runtime_is_initialized().

• Added a check for PJRT_DEVICE == "TPU" when TPU_NAME is undefined.
• Removed calls to torch_xla.device_count() before xmp.spawn because it initialised devices.

Fixes #3358

@BenjaminBossan @SunMarc @zach-huggingface

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.