PyTorch一维张量索引异常：不同维度NumPy数组索引差异原因

问题核心 #

你遇到的差异源于PyTorch对float类型numpy数组索引的特殊处理逻辑，具体与索引数组的第一个维度是否为2的幂直接相关。

详细解释 #

• 当使用形状为(32,500)的float类型numpy数组（np.zeros默认生成float64类型）索引1维整数张量时，由于32是2的5次幂，PyTorch会正确识别这是一个多维索引数组，执行高级索引操作：将数组中每个float类型的0.0隐式转换为整数0，索引张量a的第0个元素，最终返回与索引数组同形状的32×500张量。
• 当修改为形状(31,500)时，31不是2的幂，PyTorch错误地将这个多维索引数组解析为多个独立的索引维度。而被索引的张量a仅为1维，无法接受多个维度的索引，因此抛出IndexError: too many indices for tensor of dimension 1。

解决方案 #

避免用float类型数组作为整数张量的索引，显式将索引数组转换为整数类型即可解决问题：

c = np.zeros(shape=((31,500)), dtype=np.long)
print(a[c])

内容的提问来源于stack exchange，提问作者cun boss