fix: ddp for multi-node setup fixed

[bagxi]

Pull Request FAQ

• documentation
• contribution guide
• minimal examples section
• changelog for main framework updates
• Catalyst slack (#__questions channel) for issue discussion

Description

Related Issue

Type of Change

• [ ] Examples / docs / tutorials / contributors update
• [ ] Bug fix (non-breaking change which fixes an issue)
• [ ] Improvement (non-breaking change which improves an existing feature)
• [ ] New feature (non-breaking change which adds functionality)
• [ ] Breaking change (fix or feature that would cause existing functionality to change)

PR review

Anyone in the community is free to review the PR once the tests have passed. If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Checklist

• [ ] Have you updated tests for the new functionality?
• [ ] Have you added your new classes/functions to the docs?
• [ ] Have you updated the CHANGELOG?
• [ ] Have you run colab minimal CI/CD with latest requirements? Please attach the notebook link.
• [ ] Have you run colab minimal CI/CD with minimal requirements? Please attach the notebook link.
• [ ] Have you checked XLA integration? Please attach the notebook link.
• [ ] Have you checked distributed XLA integration? Please attach the notebook link.