Add invariant check (bucket indices: 0, 1, ..., k-1) (#62623)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/62623
Test Plan: Imported from OSS
Reviewed By: mrshenli
Differential Revision: D30058544
Pulled By: andwgu
fbshipit-source-id: a56910f294c6a40118751eebe255b62700f42be9
