huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-20 20:09:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Addressing a lot of comments.

Browse Source
This commit is contained in:
Nicolas Patry
2023-12-15 13:06:04 +01:00
parent aa04015098
commit 6bc92e63cb
4 changed files with 33 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
candle-metal-kernels/src/lib.rs
View File

@ -597,6 +597,7 @@ pub fn call_last_softmax(
length: usize,
elements_to_sum: usize,
input: &Buffer,
input_offset: usize,
output: &Buffer,
) -> Result<(), MetalKernelError> {
let pipeline = kernels.load_pipeline(device, Source::Reduce, kernel_name)?;
@ -604,7 +605,10 @@ pub fn call_last_softmax(
encoder.wait_for_fence(&kernels.fence);
encoder.set_compute_pipeline_state(&pipeline);
set_params!(encoder, (length, elements_to_sum, input, output));
set_params!(
encoder,
(length, elements_to_sum, (input, input_offset), output)
);
let out_length = length / elements_to_sum;