Hello all,

I am curious if you guys have a better (i.e., faster, computationally more efficient) solution than me.

Situation:

I have two tensors with size A=[bs, n, k] and B=[bs, m, k], where bs: batch size, n: number of low-level labels, m: number of high-level categories, and k: hidden dimension. I want to add the elements of B to specific elements/vectors of A based on a known index. I know exactly which element of B has to be added to one or more of the elements of A. Please see a working example of how I approach this problem (disregard the batch size as it is not important for the problem) - please note that I perform these operations on multiple GPUs:

import torch

A = torch.tensor([[0, 0, 0,], [1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4]])

B = torch.tensor([[1, 1, 1], [2, 2, 2]])

indices = [0, 1, 1, 0, 0]

for i, index in enumerate(indices):

A[i] += B[index]

print(A)

Result:

tensor([[1, 1, 1],

[3, 3, 3],

[4, 4, 4],

[4, 4, 4],

[5, 5, 5]])

Is there a more efficient way to this?

Thank you,

Christoph