[ROCm] Enable BFloat16 type for TopK operator on ROCm. (#34849)
Summary:
This PR enables bfloat16 for topk on ROCm.
iotamudelta ezyang
Pull Request resolved: https://github.com/pytorch/pytorch/pull/34849
Differential Revision: D20544732
Pulled By: ezyang
fbshipit-source-id: 1ad017a4403d2a429d98e60c8eb1f78b320df920
