Fonction torch.gather de PyTorch

Pytorch torch 参考手册Manuel de référence Pytorch torch

torch.gatherest une fonction dans PyTorch utilisée pour collecter des éléments aux indices spécifiés le long de la dimension spécifiée.

Définition de fonction

torch.gather(input, dim, index, sparse_grad)

Exemple d'utilisation

Exemple

import torch

x = torch.tensor([[1, 2], [3, 4], [5, 6]])

# Collecter le long de dim=1
index = torch.tensor([[0], [1], [0]])
result = torch.gather(x, dim=1, index=index)

print(result)

La sortie est :

tensor([[1],
        [4],
        [5]])

Pytorch torch 参考手册Manuel de référence Pytorch torch

Autres extensions