tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update readme and fmt

Browse Source
This commit is contained in:
jhqxxx
2025-11-13 01:03:52 +08:00
parent b5cfcb4684
commit 4e7573e4ff
6 changed files with 24 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/utils/mod.rs
+1 -1
View File
@@ -21,7 +21,7 @@ pub fn get_device(device: Option<&Device>) -> Device {
None => {
#[cfg(feature = "cuda")]
{
Device::new_cuda(6).unwrap_or(Device::Cpu)
Device::new_cuda(0).unwrap_or(Device::Cpu)
}
#[cfg(not(feature = "cuda"))]
{
src/utils/tensor_utils.rs
+4 -1
View File
@@ -15,7 +15,10 @@ pub fn prepare_causal_attention_mask(
let arange = Tensor::arange(0u32, tgt_len as u32, device)?;
let arange = arange.unsqueeze(1)?.broadcast_as((tgt_len, tgt_len))?;
let upper_triangle = arange.t()?.gt(&arange)?;
let mask = upper_triangle.where_cond(&Tensor::new(f32::NEG_INFINITY, device)?.broadcast_as(arange.shape())?, &Tensor::new(0f32, device)?.broadcast_as(arange.shape())?)?;
let mask = upper_triangle.where_cond(
&Tensor::new(f32::NEG_INFINITY, device)?.broadcast_as(arange.shape())?,
&Tensor::new(0f32, device)?.broadcast_as(arange.shape())?,
)?;
let mask = if seqlen_offset > 0 {
let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, device)?;