Torch sum a tensor along an axis

python sum pytorch torch tensor

Abhishek Bhatia · Jun 28, 2017 · Viewed 72.2k times · Source

ipdb> outputs.size()
torch.Size([10, 100])
ipdb> print sum(outputs,0).size(),sum(outputs,1).size(),sum(outputs,2).size()
(100L,) (100L,) (100L,)

How do I sum over the columns instead?

Answer

The simplest and best solution is to use torch.sum().

To sum all elements of a tensor:

torch.sum(outputs) # gives back a scalar

To sum over all rows (i.e. for each column):

torch.sum(outputs, dim=0) # size = [1, ncol]

To sum over all columns (i.e. for each row):

torch.sum(outputs, dim=1) # size = [nrow, 1]