[XLA:GPU] Minor improvements to the CUB sort test.
[XLA:GPU] Minor improvements to the CUB sort test.
- Add a comment
- Verify that the rewrite is actually happening and CUB sort is indeed used. I got bitten by a passing test with a temporary change that didn't actually use CUB.
- Make the test size medium. Otherwise it timed out quite a lot for me.
