I have a sequence of L events stored in a tensor of shape L x W, where for each event its last feature is a timestamp which indicates when the event appears (it is actually a musical sequence in which events are notes, and the timestamp indicates when notes are played).
What I’d like to do is then to split this sequence into blocks according to their timestamp.
More concretely, let’s assume my sequence is:
tensor([
    [n1, 0],
    [n2, 0.25],
    [n3, 0.75],
    [n4, 1],
    [n5, 1.5],
    [n6, 2.1],
    ...
])

What I want is to split the sequence into blocks so that each block contains the events appearing during a 1 second window, i.e. in the previous example return:
[
    tensor([
        [n1, 0],
        [n2, 0.25],
        [n3, 0.75],
    ]),
    tensor([
        [n4, 1],
        [n5, 1.5]
    ]),
    tensor([
        [n6, 2.1]
    ]),
    ...
]

The only solution I see is to “naively” use Python lists and for loops, so I was wondering if many exists a magic function that enables to do this in a fast way. I was thinking about using torch.split which is much faster and might do the job but I can’t find how to use it.
benchmark comparison between for loop and torch.split:
$ python3 -m timeit -s 'import torch; seq=torch.tensor([0, 0.25, 0.75, 1, 1.5, 2.1, 3]).reshape(-1,1)' '
blocks = []
block = []
limit = 1
for event in seq:
    if event[-1] >= limit:
        blocks.append(torch.stack(block, dim=0))
        limit += 1
        block = [event]
    else:
        block.append(event)
blocks.append(torch.stack(block, dim=0))
'
1000 loops, best of 3: 418 usec per loop
$ python3 -m timeit -s 'import torch; seq=torch.tensor([0, 0.25, 0.75, 1, 1.5, 2.1, 3]).reshape(-1,1)' 'torch.split(seq, [3, 2, 1, 1])'
100000 loops, best of 3: 12.2 usec per loop

Is there an efficient way to solve this problem?
Best,
Alain