PyTorch uint8张量索引报错咨询：为何与int32索引表现不同？

问题现象 #

当使用uint8类型的in_inds对同类型的signal_matrix进行索引时，触发如下错误：

IndexError: The shape of the mask [3, 3] at index 0 does not match the shape of the indexed tensor [9, 4] at index 0

但将in_inds改为int32类型后，就能得到预期的三维张量结果。

原代码 #

import torch

signal_matrix = torch.tensor(
                       [[0, 0, 1, 1],
                        [0, 0, 0, 0],
                        [0, 0, 1, 1],
                        [0, 1, 1, 0],
                        [1, 0, 1, 0],
                        [0, 0, 0, 0],
                        [0, 0, 0, 0],
                        [0, 0, 0, 0],
                        [0, 0, 0, 0]], dtype=torch.uint8)

in_inds = torch.tensor(
         [[ 0,  2,  3],                     
          [ 1,  2,  4],
          [ 0,  0,  0]][::-1], dtype= torch.uint8
          )

out_inds = torch.tensor(
         [ 5, 6, 7], dtype=torch.uint8)

op_inds = torch.tensor(
        [ [1, 1, 1, 1],                    
          [0, 0, 0, 0],
          [2, 2, 2, 2]], dtype=torch.uint8)

in_signals = signal_matrix[in_inds]

预期结果（in_inds改为int32时） #

tensor([[[0, 0, 1, 1],
         [0, 0, 1, 1],
         [0, 0, 1, 1]],

        [[0, 0, 0, 0],
         [0, 0, 1, 1],
         [1, 0, 1, 0]],

        [[0, 0, 1, 1],
         [0, 0, 1, 1],
         [0, 1, 1, 0]]], dtype=torch.int32)

uint8索引的工作机制 #

PyTorch对uint8类型的索引张量有特殊处理逻辑：

• uint8张量会被默认当作布尔掩码使用：因为布尔类型在PyTorch底层是用uint8存储的（0代表False，1代表True），所以当你传入uint8类型的张量作为索引时，PyTorch会自动将其解析为布尔掩码，而非整数索引。
• 布尔掩码的shape要求：作为掩码索引时，掩码的shape必须与被索引张量的对应维度shape匹配，或者能广播到对应维度的shape。

回到你的问题：
in_inds是shape为[3,3]的uint8张量，PyTorch将其视为布尔掩码，尝试去索引signal_matrix的第0维度（长度为9）。但[3,3]的shape无法匹配或广播到9，因此触发了IndexError。

而当in_inds改为int32等整数类型时，PyTorch会将其当作高级整数索引处理：根据in_inds中的每个整数值，从signal_matrix的第0维度取出对应行，最终拼接成与in_indsshape一致的三维张量，这就是你得到的预期结果。

解决方法 #

如果需要用uint8类型的张量作为整数索引，只需显式将其转换为整数类型即可，比如：

in_signals = signal_matrix[in_inds.int()]
# 或者指定具体整数类型
in_signals = signal_matrix[in_inds.to(torch.int32)]

内容的提问来源于stack exchange，提问作者Nobilis Manius