TORCH.GATHER
torch.gather(input, dim, index, *, sparse_grad=False, out=None) → [Tensor]Gathers values along an axis specified by dim.Parametersinput (Tensor) – the source tensordim (int) – the axis along which to indexindex (LongTensor) – the indices of elements to gatherKeyword Argumentssparse_grad (bool, optional) – If True, gradient w.r.t. input will be a sparse tensor.out (Tensor, optional) ...