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.

Best

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 http://pytorch.org/docs/0.3.1/torch.html?highlight=gather#torch.gather