How do i index a pytorch tensor?

Kong: <p>I have a simple example that is not working. I want to extract a tensor from tr_input using the indices from tr_ind.</p>
<pre><code class="lang-auto">    print(np.shape(tr_ind))
    print(np.shape(tr_input))
    
    tr_input1 = tr_input[tr_ind]
</code></pre>
<pre><code class="lang-auto">
torch.Size([150])
torch.Size([1515, 200, 28])
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
&lt;ipython-input-28-5b55de0850db&gt; in &lt;module&gt;()
      6     print(np.shape(tr_input))
      7 
----&gt; 8     tr_input1 = tr_input[tr_ind,:,:]
      9     #tr_target1 = tr_target[tr_ind]
     10     #val_input = tr_input[val_ind]

~\Anaconda3\envs\dl\lib\site-packages\torch\autograd\variable.py in __getitem__(self, key)
     76                 return IndexSelect.apply(self, 0, key)
     77             # else fall through and raise an error in Index
---&gt; 78         return Index.apply(self, key)
     79 
     80     def __setitem__(self, key, value):

~\Anaconda3\envs\dl\lib\site-packages\torch\autograd\_functions\tensor.py in forward(ctx, i, index)
     87             result = i.index(ctx.index)
     88         else:
---&gt; 89             result = i.index(ctx.index)
     90             ctx.mark_shared_storage((i, result))
     91         return result

TypeError: Performing basic indexing on a tensor and encountered an error indexing dim 0 with an object of type torch.IntTensor. The only supported types are integers, slices, numpy scalars, or if indexing with a torch.LongTensor or torch.ByteTensor only a single Tensor may be passed.
</code></pre>

miguelvr: <pre><code class="lang-auto">import torch

a = torch.rand(10, 10, 5)
print a[0, :, :]
</code></pre>
<p>this works… check the types of <code>tr_ind</code>, that might be the problem</p>
<p>PS: I’m using pytorch v0.4</p>

Kong: <p>It is a torch int tensor with 150 elements</p>

justusschock: <p>There is also the function <a href="https://pytorch.org/docs/master/torch.html#torch.index_select" rel="nofollow noopener">index_select</a> which should do exactly what you want</p>

Brando_Miranda: <p>Here is a solution if you want to index a tensor in an arbitrary dimension and select a set of tensors from that dimension (an example is say we want to compute some average of the first 3 layers):</p>
<pre><code class="lang-auto"># selecting indices arbitrarily i.e. x[*,indicies,*] were * denotes that the rest of the layers are kept the same
# but for only the last 3 layers [T, L] -&gt; [1]
x = torch.randn(5, 4)
# compute average of first 3 layer
L = x.size(1)
indices = torch.tensor(range(L-1))
xx = torch.index_select(x, dim=1, index=indices)
print(f'{x=}')
print(f'{xx=}')
print(xx.size())
x=tensor([[-0.1588, -1.2282, -0.3142, -0.0458],
        [-0.3095, -1.1847,  1.4570, -0.1697],
        [ 0.8444,  0.6528,  0.5663,  2.4564],
        [ 0.3549,  1.8841, -0.6826,  0.0283],
        [ 0.4958,  0.0893,  2.1306, -0.4176]])
xx=tensor([[-0.1588, -1.2282, -0.3142],
        [-0.3095, -1.1847,  1.4570],
        [ 0.8444,  0.6528,  0.5663],
        [ 0.3549,  1.8841, -0.6826],
        [ 0.4958,  0.0893,  2.1306]])
torch.Size([5, 3])
</code></pre>
<p>the docs are fantastically clean though: <a href="https://pytorch.org/docs/master/generated/torch.index_select.html#torch.index_select" class="inline-onebox" rel="noopener nofollow ugc">torch.index_select — PyTorch master documentation</a></p>
<hr>
<p>actually you can also do:</p>
<pre><code class="lang-auto">x.index_select(dim, index)
</code></pre>