Skip to content

Commit 0661e62

Browse files
juntaoclaude
andcommitted
Fix where_self argument order in attention mask construction
In tch-rs, `where_self(condition, other)` wraps `torch.where(condition, self, other)` where `condition` must be a Bool tensor. The previous code incorrectly passed `allow_mask` (Bool) as `self` and `zero` (Float) as `condition`, which swaps the roles of the condition and value tensors. The comment `// where(allow, 0, -inf)` makes the intended semantics clear: return 0 where allow is true, -inf where false. The correct call is `zero.where_self(&allow_mask, &neg_inf)` so that `allow_mask` (Bool) is the condition, `zero` is the value where true, and `neg_inf` is the value where false. Signed-off-by: Michael Yuan <michael@secondstate.io> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent eed686e commit 0661e62

1 file changed

Lines changed: 1 addition & 1 deletion

File tree

src/audio_encoder.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ impl AudioEncoder {
236236
);
237237
// where(allow, 0, -inf)
238238
let mask = Tensor::from_tch(
239-
allow_mask.as_tch().where_self(&zero.into_tch(), &neg_inf.into_tch())
239+
zero.into_tch().where_self(&allow_mask.as_tch(), &neg_inf.into_tch())
240240
);
241241
Some(mask)
242242
}

0 commit comments

Comments
 (0)