Advanced Selection from Tensors in Pytorch
Using torch.index_select, torch.gather and torch.take
Published in
6 min readFeb 27, 2024
In some situations, you’ll need to do some advanced indexing / selection with Pytorch, e.g. answer the question: “how can I select elements from Tensor A following the indices specified in Tensor B?”
In this post we’ll present the three most common methods for such tasks, namely torch.index_select, torch.gather and torch.take. We’ll explain all of them in detail and contrast them with one another.