[ROCm] Enable sort operator BF16 support (#71226)
Summary:
Related to [https://github.com/pytorch/pytorch/issues/58196](https://github.com/pytorch/pytorch/pull/58196)
cc jeffdaily sunway513 jithunnair-amd ROCmSupport KyleCZH amathews-amd
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71226
Reviewed By: malfet
Differential Revision: D34152115
Pulled By: seemethere
fbshipit-source-id: 53841c91976bdb5a0002362f22a54ec23aa2f78f
(cherry picked from commit 963027c7f28cf20e1c4e5722eb62b5629e735a8e)
