Solved by wetfire in post #4 
              
              
                Thanks to tom, I found out that scatter_add exists, and from there I found torch_scatter’s segment_coo 
Here’s my solution now: 
    initial_reps_list = []
    for i, sample_output in enumerate(sequence_output):
        token_mapping = all_token_mapping[i]
        token_mapping = token_mapping[token…