Adapt code for PyTorch and JAX

[buildwithsuhana]

Update DistributedBackend to respect explicit backend selection and fix PyTorch device reporting.

This change allows users to explicitly specify 'pytorch' or 'jax' as the backend, preventing auto-detection from overriding the choice. It also corrects an issue where PyTorch device information was inaccurately reported, now listing actual CUDA devices or falling back to cpu:0.

---

[cursor]

Cursor Agent can help with this pull request. Just @cursor in comments and I'll start working on changes in this branch.
_{Learn more about Cursor Agents}