PyTorch技术问题：如何用张量a的非零元素更新张量b

首先明确输入的两个张量：

import torch
a = torch.tensor([1, 2, 3, 0, 0, 1])
b = torch.tensor([0, 1, 3, 3, 0, 0])

要实现用a中的非零元素覆盖b对应位置的值，最直接高效的方式是利用PyTorch的布尔索引操作，步骤如下：

1. 生成标记a中非零位置的掩码张量：

mask = a != 0

这个掩码是和a同形状的布尔张量，a中非零的位置对应True，零位置对应False。

2. 通过掩码直接更新b的对应位置：

b[mask] = a[mask]

执行后查看结果：

print(b)  # 输出: tensor([1, 2, 3, 3, 0, 1])

这种方法基于PyTorch原生张量操作，没有Python层面的循环，能充分利用GPU加速（若使用CUDA张量），处理大规模张量时效率优势显著。

内容的提问来源于stack exchange，提问作者Ahmad