[torch][segment_reduce] Add support for int lengths as well (#61141)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61141
Currently only long is supported. This diff adds support for other index type.
Next Steps:
- Update default, refactor unit test and test non_initial value as well
- Cleanup (more tests, benchmark, update documentation)
Test Plan: updated unit test. rely on CI.
Reviewed By: ngimel
Differential Revision: D29526308
fbshipit-source-id: b4043603483851ef7e0e93b0bb02ac7849c6449d
