xla
Use `_sharded_cpu_state_dict` functionality to Write Items for SPMD Save Planner
#5315
Merged
Merged
Use `_sharded_cpu_state_dict` functionality to Write Items for SPMD Save Planner #5315
shahyash10
merged
15 commits into
master
from
yashjs_use_sharded_cpu_dict
initial commit
337f6970
add suggested changes
eeeff609
jonb377
commented
on 2023-07-18
add unit test
65c87e6a
fix test
a890aa73
fix test
eab0046d
add suggested changes
1fcf7574
remove is_sharded_tensor check
259ceb66
check if device type is xla in `wrap_if_sharded`
f7faaf60
jonb377
commented
on 2023-07-18
change order
d86a5076
jonb377
commented
on 2023-07-18
update resolve_data and add more tests
f7daeb7d
Merge branch 'master' of github.com:pytorch/xla into yashjs_use_shard…
ad8b8058
run linter
c34bb850
use subtest
5ffa4760
formatting fixes
4124209e
run linter
488b5ac6
Login to write a write a comment.
Login via GitHub