How to implement this (essentially calculating average of vectors)
batch_size, d= 8, 128
n, N= 100, 10

feature=torch.rand(batch_size, n, d)
belong=torch.randint(0, N, (batch_size,n))
output=torch.zeros(batch_size, N, d)

for i in range(batch_size):
    for j in range(N):
        output[i][j]=feature[i][belong[i]==j].mean(0)

with high-performance, CUDA-friendly and autograd-friendly pytorch function?
Grateful to your help 
P.S. I tried torch.index_add and it performs badly for CUDA