`sort_last_dim` fails on cuda

[lucasavila00]

trafficstars

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

thread '<unnamed>' panicked at /src/lib.rs:133:25:
run_engine error: DriverError(CUDA_ERROR_INVALID_VALUE, "invalid argument")
   0: <core::result::Result<O,E> as candle_core::cuda_backend::error::WrapErr<O>>::w
   1: <candle_core::sort::ArgSort as candle_core::custom_op::CustomOp1>::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