pytorch

6af5d00e - [torch][segment_reduce] Add support for multi-dimensional input (cuda) (#60018)

[torch][segment_reduce] Add support for multi-dimensional input (cuda) (#60018) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60018 Same as title. This diff finishes cuda support for currently implemented reductions and input parameters. Next Steps: - Add support for sum/min - More testing and benchmarking - Cleanup - Update default values when length is 0 - Use TensorIterator - Update documentation Test Plan: Unit test to cover cuda forward path. Reviewed By: ngimel Differential Revision: D29135373 fbshipit-source-id: d070727eeb660f56782e7ac8a5b0798be688480a