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 range(batch_size):
    indices.append(torch.randint(high=size, size=(n_idx**2, 2)))

def forlooping(S, indices):
    # Mock indices
    sub_S = list()
    for b in range(batch_size):
        sub_S.append(S[b][indices[b][:, 0], indices[b][:, 1]].reshape(n_idx, n_idx))
    sub_S = torch.stack(sub_S)
    return sub_S

def indexing(S, indices):
    indices = torch.stack(indices, dim=0).view(-1, n_idx, n_idx, 2)
    bidx = torch.arange(batch_size).view(-1,1,1)
    sub_S = S[bidx, indices[...,0], indices[...,1]]
    return sub_S

The speed up is noticeable with large batch sizes like so
In [2]: %timeit indexing(S, indices)
155 µs ± 1.98 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [3]: %timeit forlooping(S, indices)
1.9 ms ± 14.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

To eliminate the first for loop, simply include the batch size in the size parameter like below
indices = torch.randint(high=size, size=(batch_size, n_idx**2, 2)))

Hope this helps!