Solved by Mazhar_Shaikh in post #2 
              
              
                Hi Samue1, 
To eliminate the second for loop in your code, you can use broadcasted indexes in the batch dimension as well. An example is shown below 
batch_size = 100
size = 24
n_idx = 4

# Mock data
S = torch.arange(end=batch_size*size*size).reshape(batch_size, size, size)
indices = list()
for _ in…