1. torch.gather(input, dim, index)
import torch
a = torch.Tensor([[1,2],[3,4]])
>>> a
1 2
3 4
b = torch.gather(a,1,torch.LongTensor([[0,0],[1,0]]))
>>> b
1 1
4 3
index和input维度一样,存储的是index元素的下标。
a = torch.Tensor([[1,2],[3,4]]) b = torch.gather(a,1,torch.LongTensor([[0,0],[1,0]])) 其中, dim=1,表示的是在第二维度上操作。
index = torch.LongTensor([[0,0],[1,0]]),[0,0]就是第一行对应元素的下标,也就是对应的是[1,1]; [1,0]就是第二行对应元素的下标,也就是对应的是[4,3]。
index类型貌似必须为LongTensor。