Refactor BucketBatch (#63185)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/63185
Test Plan: Imported from OSS
Reviewed By: bdhirsh
Differential Revision: D30288893
Pulled By: ejguan
fbshipit-source-id: b88b792d12a83c99d8ea9e516e3b4c54a82100f6
