I've got the same issue. For what its worth, casting to float16 works as well. The issue also only occurs when the tensor is on the GPU. If you move the tensor to the CPU, then sort, sorting works as expected.