xla
Allow MpDeviceLoader to shard dictionaries of tensors
#7912

Merged

Allow MpDeviceLoader to shard dictionaries of tensors #7912

bhavya01 merged 4 commits into master from mp_loader

Allow MpDeviceLoader to shard dictionaries of tensors

55fedb4a

Lint fix

b9d98aef

bhavya01 marked this pull request as ready for review 1 year ago

bhavya01 requested a review from

JackCaoG 1 year ago

bhavya01 requested a review from

will-cromar 1 year ago

bhavya01 assigned

bhavya01 1 year ago

Add test_mp_input_sharding.py to run_test.sh

c24b6e06

JackCaoG commented on 2024-08-26

bhavya01 added tpuci

Update send_data_to_cpu to allow none input_sharding when batch is a …

bb23390b

JackCaoG approved these changes on 2024-09-04

bhavya01 merged 5acecc33 into master 1 year ago

Reviewers

JackCaoG

will-cromar

Assignees

bhavya01

Labels

None yet

Milestone

No milestone