From 6c4a960b15404b9307328fa4e2c929f813b6b092 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 27 Jun 2023 09:56:19 +0100 Subject: [PATCH] Embedding bugfix. --- examples/cuda_basics.rs | 4 ++-- kernels/src/embeddings.cu | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/cuda_basics.rs b/examples/cuda_basics.rs index 2c790ed1..c8852bf6 100644 --- a/examples/cuda_basics.rs +++ b/examples/cuda_basics.rs @@ -3,8 +3,8 @@ use candle::{Device, Tensor}; fn main() -> Result<()> { let device = Device::new_cuda(0)?; - let ids = Tensor::new(&[0u32, 3u32, 1u32], &device)?; - let t = Tensor::new(&[[0f32, 1f32], [1f32, 2f32], [2f32, 3f32]], &device)?; + let ids = Tensor::new(&[0u32, 2u32, 1u32], &device)?; + let t = Tensor::new(&[[0f32, 1f32], [2f32, 3f32], [4f32, 5f32]], &device)?; let hs = Tensor::embedding(&ids, &t)?; println!("> {:?}", hs.to_vec2::()); diff --git a/kernels/src/embeddings.cu b/kernels/src/embeddings.cu index 1dd12cf1..8f27b0c9 100644 --- a/kernels/src/embeddings.cu +++ b/kernels/src/embeddings.cu @@ -18,13 +18,13 @@ extern "C" __global__ void FN_NAME( \ const size_t *strides = info + num_dims; \ if (is_contiguous(num_dims, dims, strides)) { \ for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ - memcpy(out + i * h_size, inp + ids[i], h_size); \ + memcpy(&out[i * h_size], &inp[ids[i] * h_size], h_size * sizeof(TYPENAME)); \ } \ } \ else { \ for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \ - memcpy(out + i * h_size, inp + ids[i], h_size); \ + memcpy(&out[i * h_size], &inp[ids[strided_i] * h_size], h_size * sizeof(TYPENAME)); \ } \ } \ } \