pytorch笔记

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。