xla
Use `_sharded_cpu_state_dict` functionality to Write Items for SPMD Save Planner
#5315
Merged

Commits
  • initial commit
    shahyash10 committed 2 years ago
  • add suggested changes
    shahyash10 committed 2 years ago
  • add unit test
    shahyash10 committed 2 years ago
  • fix test
    shahyash10 committed 2 years ago
  • fix test
    shahyash10 committed 2 years ago
  • add suggested changes
    shahyash10 committed 2 years ago
  • remove is_sharded_tensor check
    shahyash10 committed 2 years ago
  • check if device type is xla in `wrap_if_sharded`
    shahyash10 committed 2 years ago
  • change order
    shahyash10 committed 2 years ago
  • update resolve_data and add more tests
    shahyash10 committed 2 years ago
  • Merge branch 'master' of github.com:pytorch/xla into yashjs_use_sharded_cpu_dict
    shahyash10 committed 2 years ago
  • run linter
    shahyash10 committed 2 years ago
  • use subtest
    shahyash10 committed 2 years ago
  • formatting fixes
    shahyash10 committed 2 years ago
  • run linter
    shahyash10 committed 2 years ago
Loading