But what I want to happen is to select certain indices for each dimension, that is if I were to do

test=torch.ones(10,4)
print(test[torch.arange(9),torch.arange(3)])

then it would print a 2-D tensor of ones of size [9,3]; that is, the first 9 rows, and for each of those rows, the first 3 elements. What would be the equivalent functionality? The above code returns

IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes [9], [3]