Solved by ptrblck in post #2 
              
              
                A combination of gather and scatter_ should work: 
x = torch.tensor([[1,2,3],
                  [4,5,6],
                  [7,8,9]])

y = torch.tensor([[11,21,31],
                  [41,51,61],
                  [71,81,91]])

z = torch.tensor([[1],
                  [2],
                  [0]])


k …