Use u8 tensors for masks. (#273)

This commit is contained in:
Laurent Mazare
2023-07-29 11:32:58 +01:00
committed by GitHub
parent 50d8273ae4
commit 4bf2ebf836
6 changed files with 6 additions and 10 deletions

View File

@ -424,7 +424,7 @@ pub struct Falcon {
fn make_causal_mask(t: usize) -> Result<Tensor> {
let mask: Vec<_> = (0..t)
.flat_map(|i| (0..t).map(move |j| u32::from(j > i)))
.flat_map(|i| (0..t).map(move |j| u8::from(j > i)))
.collect();
let mask = Tensor::from_slice(&mask, (t, t), &Device::Cpu)?;
Ok(mask)