diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index 566b5f36b..660ba65e2 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -95,8 +95,8 @@ struct CommandEncoder { return enc_->setBytes(&v, sizeof(T), idx); } - void set_threadgroup_memory_length(size_t length, NS::UInteger index) { - enc_->setThreadgroupMemoryLength(length, index); + void set_threadgroup_memory_length(size_t length, int idx) { + enc_->setThreadgroupMemoryLength(length, idx); } ConcurrentContext start_concurrent() {