This depends on pytorch/xla#6683, @wanchaol should we skip distribute_module test here, and just rely on the downstream test? I remember that you were suggesting it in the refactoring PR. If we skip here, then this doesn't need to be blocked on pytorch/xla#6683.

I am good keeping it as a safeguard here, to ensure that we can call the xla_distribute_module API via distribute_module here.

@pytorchbot merge

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Not updating the xla pin, since the latest one seems to have some pjrt gpu plugin build issue... I hope the original one in main is intact cc @vanbasten23

@pytorchbot merge

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

This is also related to pytorch/xla#6322