Torch sum each row excluding an index

Question

Given a Tensor A of shape (N,C) and an indices Tensor Idx of shape (N,), i'd like to sum all the elements of each row in A excluding the corresponding column index in I. For example:

A = torch.tensor([[1,2,3], [4,5,6]]) Idx = torch.tensor([0,2]) #result: torch.tensor([[5], [9]])

A solution using loops is known.

Anton Ganichev · Accepted Answer · 2020-04-06 20:10:11Z

You can set excluded elements to zero:

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

and sum tensor along rows:

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

Collectives™ on Stack Overflow

Torch sum each row excluding an index

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related