diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index 7d0da0580..539704fe1 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -368,7 +368,7 @@ void CommandEncoder::commit() { } void CommandEncoder::synchronize() { - cudaStreamSynchronize(stream_); + CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_)); auto p = std::make_shared>(); std::future f = p->get_future(); add_completed_handler([p = std::move(p)]() { p->set_value(); }); diff --git a/mlx/backend/cuda/worker.cpp b/mlx/backend/cuda/worker.cpp index c468b8501..495204d3d 100644 --- a/mlx/backend/cuda/worker.cpp +++ b/mlx/backend/cuda/worker.cpp @@ -44,7 +44,7 @@ void Worker::commit(cudaStream_t stream) { } signal_event_.record(stream); signal_event_.wait(signal_stream_); - cudaLaunchHostFunc(signal_stream_, signal, this); + CHECK_CUDA_ERROR(cudaLaunchHostFunc(signal_stream_, signal, this)); } void Worker::thread_fn() {