Indexing torch tensors with a list of indices

1

I populate a tensor with dimensions that are not fixed. Then, I try to index this tensor. I would like to pass the indices as a list, but that doesn’t seem to work.

For ex:

A = torch.rand((5,5,4,18))
B = [3,3,2]

Instead of indexing it as

A[3,3,2]

why can’t I do

A[B]

I need this because in my case the dimensions of B keep changing.

2

This should work:

A = torch.rand((5,5,4,18))
B = [3, 3, 2]
print((A[3, 3, 2]==A[torch.tensor(B).split(1)]).all())
> tensor(True)