Bring back old algorithm for sorting on small number of segments (#64127)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/63456
The code was copy-pasted from the previous commit without modification.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64127
Reviewed By: mruberry
Differential Revision: D30632090
Pulled By: ngimel
fbshipit-source-id: 58bbdd9b0423f01d4e65e2ec925ad9a3f88efc9b