火炬总和不包括索引的每一行

【问题标题】：Torch sum each row excluding an index火炬总和不包括索引的每一行
【发布时间】：2020-04-06 19:44:55
【问题描述】：

给定一个形状为 (N,C) 的张量 A 和一个形状为 (N,) 的索引张量 Idx，我想将 A 中每一行的所有元素相加，不包括 I 中的相应列索引。例如：

A = torch.tensor([[1,2,3],
                  [4,5,6]])

Idx = torch.tensor([0,2])

#result:
torch.tensor([[5],
              [9]])

使用循环的解决方案是已知的。

【问题讨论】：

【解决方案1】：

您可以将排除的元素设置为零：

A[range(A.shape[0]),Idx] = 0

并沿行对张量求和：

b = A.sum(dim = 1,keepdim = True ) # b = torch.tensor([[5],  [9]])

【讨论】：