torch.Tensor.take_along_dim

Tensor.take_along_dim(indices, dim) → Tensor

参见 torch.take_along_dim()