Indexed assigning 1-D Variable in 4-D Variable

Hello all,
I have a task which is indexed assigning 1-D tensor, s in a target 4-D tensor, A. 4-D tensor has size b x c x w x h, and the 1-D tensor has the length of c. For each sample(b sample total), I have a different binary mask (w x h). For each location of mask value equal to 1, I want to assign the same 1-D tensor for all samples. The worst plan I could come up is to do b loops: A[i, :, mask[i]] = s. But even so, Pytorch does not seem to support this in 0.3.1.

Another important thing to note is that I want gradients flow through s, which means using numpy to do the task will not be acceptable.


I’m having difficulty following exactly what you’re looking for. If you could provide a code snippet that does your indexing in some for loops that would be great.

I think what you’re looking for can be written using