How to do this in PyTorch?

1

a = torch.tensor([1,2,3,4,5,6,7,8,9])
b = [3,2,4]
I want to get a result == tensor([6, 9, 30]) where 6 == 1+2+3(three elements), 9 == 4+5(two elements), 30 == 6+7+8+9(four elements).
How to do this in Pytorch efficiently?

2

Not sure if best solution, but it yields the required result:

import torch
a = torch.tensor([1,2,3,4,5,6,7,8,9])
b = [3,2,4]
chunks = torch.split(a, b)
results = torch.tensor([ c.sum() for c in chunks ])