Pytorch, retrieving values from a tensor using several indices. Most computationally efficient solution
If I have an example 3d tensor
a = [[4, 2, 1, 6],[1, 2, 3, 8], [92, 4, 23, 54]]
tensor_a = torch.tensor(a)

I can get 2 of the 1D tensors along the first dimension using
tensor_a[[0, 1]]

tensor([[4, 2, 1, 6],
        [1, 2, 3, 8]])

But how about using several indices?
So I have something like this
list_indices = [[0,  0], [0,2], [1, 2]]
I could do something like
combos = []
for indi in list_indices:
    combos.append(tensor_a[indi])

But I'm wondering if since there's a for loop, if there's a more computationally way to do this, perhaps also using pytorch

It is more computationally effecient to use the predefined Pytorch function &quot;torch.index_select&quot; to select tensor elements using a list of indices:
a = [[4, 2, 1, 6],[1, 2, 3, 8], [92, 4, 23, 54]]
tensor_a = torch.tensor(a)
list_indices = [[0,  0], [0,2], [1, 2]]
#convert list_indices to Tensor
indices = torch.tensor(list_indices)
#get elements from tensor_a using indices.
tensor_a=torch.index_select(tensor_a, 0, indices.view(-1))
print(tensor_a)

if you want the result to be a list not a tensors, you can convert tensor_a to a list:
tensor_a_list = tensor_a.tolist()

To test the computational efficiency I created 1000000 indices and I compared the execution time. Using the loop takes more time then using my suggested pytorch approach:
import time
import torch
start_time = time.time()
a = [[4, 2, 1, 6],[1, 2, 3, 8], [92, 4, 23, 54]]
tensor_a = torch.tensor(a)
indices = torch.randint(0, 2, (1000000,)).tolist()
for indi in indices:
   combos.append(tensor_a[indi])
print(&quot;--- %s seconds ---&quot; % (time.time() - start_time))
--- 3.3966853618621826 seconds ---

start_time = time.time()
indices = torch.tensor(indices)
tensor_a=torch.index_select(tensor_a, 0, indices)
print(&quot;--- %s seconds ---&quot; % (time.time() - start_time))
--- 0.10641193389892578 seconds ---

