test/ops/self_attention.py 在 CUDA 上因 attention mask 位于 CPU 而失败

[Lonnyyy123]

运行：

python test/ops/self_attention.py --device nvidia

时，torch_self_attention 会因为 attention mask 和 attention bias 位于不同设备而失败。

问题代码：

attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=S-L)
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))

当 query 位于 CUDA 时，attn_bias 位于 CUDA，但 temp_mask 默认位于 CPU，导致设备不一致错误。

建议修改为：

temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=S-L)

另外，attn_bias.to(query.dtype) 不会原地修改 Tensor； attn_bias 创建时已经使用了正确的 dtype。