PyTorch - get list of sums of 2D tensors in 3D tensor

2
votes

I have a 3D tensor consisting of 2D tensors, e. g.:

t = torch.tensor([[[0, 0, 1],
                   [0, 1, 0],
                   [1, 0, 0]],

                  [[0, 0, 1],
                   [0, 1, 0],
                   [1, 0, 0]],

                  [[0, 0, 1],
                   [0, 1, 0],
                   [1, 0, 0]]
                  ])

I need a list or tensor of sums of those 2D tensors, e. g.: sums = [3, 3, 3]. So far I have:

sizes = [torch.sum(t[i]) for i in range(t.shape[0])]

I think this can be done with PyTorch only, but I've tried using torch.sum() with all possible dimensions and I always get sums over the individual fields of those 2D tensors, e. g.:

[[0, 0, 3],
[0, 3, 0],
[3, 0, 0]]

How can this be done in PyTorch?

2
pythonpytorch

2 Answers

3
votes

You can do it at once by passing dims as tuple.

t.sum(dim=(0,1))
tensor([3, 3, 3])

or for list

t.sum(dim=(0,1)).tolist()
[3, 3, 3]
1
votes

If understood your problem correctly, this should do the job:

t.sum(0).sum(1).tolist()

Output: [3, 3, 3]