Solved by ptrblck in post #2 
              
              
                The creation of z is not working using the provided shapes and you should probably use torch.stack instead. 
Anyway, this code should yield the same result: 
N = 100
y = torch.randn(N, 10, 10)
a = torch.randint(0, 10, (N, 15))
b = torch.randint(0, 10, (N, 15))

z = y[torch.arange(N).unsqueeze(1), a,…