我可以使用逻辑索引或索引列表对张量进行切片吗？

我正在尝试使用列上的逻辑索引对 PyTorch 张量进行切片。我想要与索引向量中的 1 值相对应的列。切片和逻辑索引都是可能的，但是它们可以一起吗？如果是这样，怎么办？我的尝试不断抛出无用的错误

类型错误：使用 ByteTensor 类型的对象对张量进行索引。这 仅支持的类型是整数、切片、numpy 标量和 torch.LongTensor 或 torch.ByteTensor 作为唯一的参数。

MCVE

所需输出

import torch

C = torch.LongTensor([[1, 3], [4, 6]])
# 1 3
# 4 6

逻辑索引仅在列上：

A_log = torch.ByteTensor([1, 0, 1]) # the logical index
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[:, A_log] # Throws error

如果向量大小相同，则逻辑索引有效：

B_truncated = torch.LongTensor([1, 2, 3])
C = B_truncated[A_log]

我可以通过重复逻辑索引来获得所需的结果，使其与我正在索引的张量具有相同的大小，但随后我还必须重塑输出。

C = B[A_log.repeat(2, 1)] # [torch.LongTensor of size 4]
C = C.resize_(2, 2)

我也尝试过使用索引列表:

A_idx = torch.LongTensor([0, 2]) # the index vector
C = B[:, A_idx] # Throws error

如果我想要连续的索引范围，slicing works:

C = B[:, 1:2]

我认为这是作为index_select https://pytorch.org/docs/stable/generated/torch.index_select.html#torch.index_select功能，你可以试试

import torch

A_idx = torch.LongTensor([0, 2]) # the index vector
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B.index_select(1, A_idx)
# 1 3
# 4 6