Solved by albanD in post #4 
              
              
                It’s still not 100% clear to me what this does  . But I guess that would work: 
dim = 2
new_size = inp.size()[:dim] + selector.size() + inp.size()[dim+1:]
out = inp.index_select(selector.view(-1), dim=dim)
out.view(new_size)