Hi folks!
I have been going through Pytorch documentation in search for a way to do an efficient per-index replacement of values inside a tensor. The problem:

dest - 4D tensor (N,  H, W, C1) that I want to update
idxs - 4D tensor of indexes I want to access (N, H, W, C2)
source - 4D tensor of values I want to put in z at idxs (N, H, W, C2).

In practice dest a minibatch of per-pixel probabilities, where C1 is a number of distinct categories and C2 is top 5 updates from source. Note that C1 != C2.
I have been looking at torch.copy_index_ but I can’t seem to find a way to broadcast it along N, H, W dimensions.
My (working) solution with a for loop:
def col_wise_replace(dest, idxs, source):
''' replace probability stacks '''
        for k in range(n):
            for i in range(h):
                for j in range(w):
                    dest[k,i,j].index_copy_(0, idxs[k,i,j], source[k,i,j])

This solution is slow. Any numpy / torch-like way I can achieve the same?