Fix subset of hosts dataloading for TPU v4#586

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

mohit/subset_v4AI-Hypercomputer/maxtext:mohit/subset_v4Copy head branch name to clipboard