Indexing a multi-dimensional tensor with a tensor in PyTorch

pytorch tensor
Chum-Chum Scarecrows picture Chum-Chum Scarecrows · Aug 30, 2018 · Viewed 13.8k times · Source

I have the following code:

a = torch.randint(0,10,[3,3,3,3])
b = torch.LongTensor([1,1,1,1])

I have a multi-dimensional index b and want to use it to select a single cell in a. If b wasn't a tensor, I could do:

a[1,1,1,1]

Which returns the correct cell, but:

a[b]

Doesn't work, because it just selects a[1] four times.

How can I do this? Thanks

Answer

dennlinger picture dennlinger · Aug 30, 2018

A more elegant (and simpler) solution might be to simply cast b as a tuple:

a[tuple(b)]
Out[10]: tensor(5.)

I was curious to see how this works with "regular" numpy, and found a related article explaining this quite well here.

Related questions

How does the "view" method work in PyTorch?

I am confused about the method view() in the following code snippet. class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2,2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = …

Read More
Best way to save a trained model in PyTorch?

I was looking for alternative ways to save a trained model in PyTorch. So far, I have found two alternatives. torch.save() to save a model and torch.load() to load a model. model.state_dict() to save a trained …

Read More
PyTorch: How to get the shape of a Tensor as a list of int

In numpy, V.shape gives a tuple of ints of dimensions of V. In tensorflow V.get_shape().as_list() gives a list of integers of the dimensions of V. In pytorch, V.size() gives a size object, but how …

Read More