Python PyTorch take用法及代碼示例

本文簡要介紹python語言中 torch.take 的用法。

用法:

torch.take(input, index) → Tensor

參數：

• input(Tensor) -輸入張量。

• index(LongTensor) -張量的索引

在給定索引處返回具有 input 元素的新張量。輸入張量被視為被視為一維張量。結果與索引的形狀相同。

例子：

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