Solved by futscdav in post #2 
              
              
                You have to expand the dims of the indexing tensor for gather, or at least I’m not aware of a way to circumvent that. This should work 
a = torch.rand((264, 2, 3))
b = torch.randint(0, 2, (264,))
i = torch.stack(3*[b], dim=1).view(-1, 1, 3)
r = a.gather(1, i).squeeze()