[WIP][DataLoader] Implement BucketBatchIterableDataset (#51126)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51126
BucketBatch:
Get a chunk of data as a bucket, and sort the bucket by the specified key, then batching.
If sort key is not specified, directly use batchIterableDS..
1. Implement BucketBatch for bucket sampler
2. Improve BatchDS tests
Test Plan: Imported from OSS
Reviewed By: H-Huang
Differential Revision: D26209890
Pulled By: ejguan
fbshipit-source-id: 8519e2e49da158b3fe32913c8f3cadfa6f3ff1fc
