torch.take

torch.take(input, index) → Tensor

傳回一個新張量，其中包含 input 中指定索引的元素。輸入張量會被視為如同它被視為一個一維張量。結果採用與索引相同的形狀。

參數

• input (Tensor) – 輸入張量。

• index (LongTensor) – 張量的索引

範例

>>> src = torch.tensor([[4, 3, 5],
...                     [6, 7, 8]])
>>> torch.take(src, torch.tensor([0, 2, 5]))
tensor([ 4,  5,  8])