Thanks for the reply @albanD . I naively used for loop to find max value and assign it to the table and
repeat this process all data point in one batch
I am not familiar with gather() and scatter() function, so could you give more hints?
def scatter_max(attn_scores, indices):
    tables = torch.zeros_like(attn_scores)
    for score, idx in zip(attn_scores, indices):
        max_score = tables[idx]
        if score > max_score:
            tables[idx] = score
    return tables