How to set a 3-d tensor to 0 according to the value of a 2-d tensor
Suppose I have a 3-d tensor P shaped (B, N, d) and a 2-d tensor Q shaped (B,N), where values in Q are smaller than d. I want to set some values in P to 0 using the indices from Q:
For instance,
P = torch.randn(2,3,6)
Q = torch.tensor([[0,3,4], [2,1,3]])

How can I set P[0,0,0]=0; P[0,1,3]=0; P[0,2,4]=0; P[1,0,2]=0; P[1,1,1]=0; P[1,2,3]=0 using Q and keep other values in P unchanged?

What you are looking to do is:
P[i, j, Q[i, j]] = 0

This is the perfect use case for torch.Tensor.scatter, which has the effect of placing an arbitrary value at designated positions.
We first need the input and indexer to have matching shapes:
&gt;&gt;&gt; Q_ = Q[...,None].expand_as(P)

And apply the scatter function on dim=2, with the hidden value argument...
P.scatter(dim=2, index=Q_, value=0)
tensor([[[0, 2, 5, 4, 3, 7],
         [4, 1, 5, 0, 3, 8],
         [7, 8, 4, 3, 0, 1]],

        [[7, 1, 0, 1, 1, 3],
         [5, 0, 5, 3, 0, 5],
         [6, 3, 5, 0, 7, 2]]])

