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])