huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-18 19:47:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: fix index_select cuda kernel for src target dim different than ids dim when selecting dim > 0 (#1037)

Browse Source 
* fix: fix index_select cuda kernel for src target dim different than ids dim when selecting dim > 0

* cargo fmt
This commit is contained in:
Gonzalo
2023-10-05 14:46:13 -03:00
committed by GitHub
parent f0c619a4af
commit 8f7973958c
3 changed files with 21 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
candle-core/src/cuda_backend.rs
View File

@ -891,7 +891,8 @@ impl<'a> Map1 for IndexSelect<'a> {
};
let left_size: usize = src_l.dims()[..self.2].iter().product();
let right_size: usize = src_l.dims()[self.2 + 1..].iter().product();
let dim_size = ids_shape.elem_count();
let src_dim_size = src_l.dims()[self.2];
let ids_dim_size = ids_shape.elem_count();
let dst_el = ids_shape.elem_count() * left_size * right_size;
let cfg = LaunchConfig::for_num_elems(dst_el as u32);
let func = dev.get_or_load_func(&kernel_name::<T>(name), kernels::INDEXING)?;
@ -905,7 +906,8 @@ impl<'a> Map1 for IndexSelect<'a> {
&src,
&out,
left_size,
dim_size,
src_dim_size,
ids_dim_size,
right_size,
);
// SAFETY: ffi.