Факел сумма каждой строки, исключая индекс

Question

Учитывая Tensor A формы (N, C) и индекс Tensor Idx формы (N,), я хотел бы суммировать все элементы каждой строки в A, исключая соответствующий индекс столбца в I. Например:

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

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

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

Решение с использованием циклов известно.

Answer 1

Вы можете установить исключенные элементы в ноль:

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

и суммировать тензор по строкам:

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