Solved by ptrblck in post #2 
              
              
                I assume ret should be initialized with the shapes [dim1, dim2, dim4], otherwise you’ll get a shape mismatch error in the posted loop. 
If that’s the case, this code should work: 
dim1 = 3
dim2 = 4
dim3 = 5
dim4 = 10
source = torch.randn((dim1,dim2,dim3,dim4))
index = torch.randint(dim3, (dim1, dim2…