huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +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

3
candle-nn/src/ops.rs
View File

@ -220,7 +220,7 @@ impl candle::CustomOp1 for SoftmaxLastDim {
};
let n = layout.stride().len();
if !(layout.is_contiguous() && layout.stride()[n - 1] == 1 && layout.start_offset() == 0) {
if !(layout.is_contiguous() && layout.stride()[n - 1] == 1) {
candle::bail!("Non contiguous softmax-last-dim is not implemented");
}
@ -235,6 +235,7 @@ impl candle::CustomOp1 for SoftmaxLastDim {
elem_count,
last_dim,
storage.buffer(),
layout.start_offset() * storage.dtype().size_in_bytes(),
&mut output,
)
.unwrap();