Search code examples
pythonpython-3.xpytorchtensortensor-indexing

Slicing uneven columns from tensor array


I have an array like so:

([[[ 0,  1,  2],
 [ 3,  4,  5]],

[[ 6,  7,  8],
[ 9, 10, 11]],

[[12, 13, 14],
[15, 16, 17]]])

If i want to slice the numbers 12 to 17 i would use:

arr[2, 0:2, 0:3]

but how would i go about slicing the array to get 12 to 16?


Solution

  • You'll need to "flatten" the last two dimensions first. Only then will you be able to extract the elements you want:

    xf = x.view(x.size(0), -1)  # flatten the last dimensions
    xf[2, 0:5]
    
    Out[87]: tensor([12, 13, 14, 15, 16])