Solved by ptrblck in post #2 
              
              
                You could flatten the tensors and use gather to get the desired output: 
out = torch.gather(torch.flatten(m, 2), 2, torch.flatten(idx, 2)).view(idx.size())
print(out)
>  tensor([[[[15., 16.],
            [19., 20.]],

           [[26., 29.],
            [38., 41.]]]])