torch.Tensor.take_along_dim

Tensor.take_along_dim(indices, dim) → Tensor

See torch.take_along_dim()