Hi there,

This is my first question, I hope that I’m in the right category!

Briefly, I have a 2d boolean mask, and I want to obtain the sum at each position.

Example:

mask = torch.Tensor([[1, 0, 0, 1, 1, 0, 1],[0, 1, 0, 1, 0, 1, 0]])

A normal sum on axis 1 will be:

torch.sum(mask, dim=1)

tensor([4., 3.])

But I would like to obtain the sum at each position, where the result has the same shape as mask:

sum_pos(mask, dim=1)

tensor([[1, 1, 1, 2, 3, 3, 4],[0, 1, 1, 2, 2, 3, 3]])

Is there any way to do that, **without a for loop on dim=1** ?

Thanks in advance!