`sort_last_dim` fails on cuda

```
[mlx_core/src/sampler.rs:48:9] &probs = Tensor[dims 32000; f32, cuda:0]

thread '' panicked at /src/lib.rs:133:25:
run_engine error: DriverError(CUDA_ERROR_INVALID_VALUE, "invalid argument")
0: as candle_core::cuda_backend::error::WrapErr>::w
1: ::cuda_fwd
2: candle_core::storage::Storage::apply_op1
---

| NVIDIA-SMI 545.23.08 Driver Version: 545.23.08 CUDA Version: 12.3 |
| 0 NVIDIA GeForce RTX 2070 On | 00000000:09:00.0 On | N/A |

```

Adding a ` .to_device(&Device::Cpu)?` before it fixes it and produces correct results.

I'm using the latest commit `13c64f6828360a9cb9b58b4f817e4f3b8316388c`