ptrblck:

If so, did a scatter operation work for you or were you hitting any issues with it?


Thank you so much @ptrblck !
The segment_csr function from torch_scatter is just exactly what I wanted.
from torch_tensor import segment_csr
x = torch.tensor([[6, 0, 3],
       [0, 7, 0],
       [4, 6, 8],
       [1, 6, 7],
       [2, 6, 1],
       [8, 5, 5],
       [9, 8, 3],
       [4, 9, 1],
       [5, 8, 6],
       [6, 6, 3]])
splits = torch.tensor([0,2,3,6,7,10])
segment_csr(x, splits,reduce='max')
# tensor([[6, 7, 3],
        [4, 6, 8],
        [8, 6, 7],
        [9, 8, 3],
        [6, 9, 6]])
# Exact answer
Thank you once again @ptrblck 
Thank you to @J_Johnson  as well