huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 11:08:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make the r, k, v tensors contiguous. (#1719)

Browse Source
This commit is contained in:
Laurent Mazare
2024-02-16 09:17:35 +01:00
committed by GitHub
parent 7c7400fb63
commit 5ebcfeaf0f
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
candle-transformers/src/models/rwkv_v5.rs
View File

@ -165,9 +165,9 @@ impl SelfAttention {
let mut out: Vec<Tensor> = Vec::with_capacity(t); let mut out: Vec<Tensor> = Vec::with_capacity(t);
for t_ in 0..t { for t_ in 0..t {
// //
let rt = receptance.i((.., .., t_..t_ + 1))?; let rt = receptance.i((.., .., t_..t_ + 1))?.contiguous()?;
let kt = key.i((.., .., .., t_..t_ + 1))?; let kt = key.i((.., .., .., t_..t_ + 1))?.contiguous()?;
let vt = value.i((.., .., t_..t_ + 1))?; let vt = value.i((.., .., t_..t_ + 1))?.contiguous()?;
let at = kt.matmul(&vt)?; let at = kt.matmul(&vt)?;
let rhs = (time_faaaa.broadcast_mul(&at)? + &state_)?; let rhs = (time_faaaa.broadcast_mul(&at)? + &state_)?;
let out_ = rt.matmul(&rhs)?.squeeze(2)?; let out_ = rt.matmul(&rhs)?.squeeze(2)?;