PyTorch张量中布尔掩码与常规索引组合索引的实现问题

这是个典型的PyTorch和NumPy索引行为差异导致的问题，我来帮你理清原因和解决办法：

为什么会报错？ #

PyTorch的布尔索引规则比NumPy更严格：当你混合布尔掩码和其他索引方式时，掩码的形状必须和被索引张量的对应维度完全匹配，或者能正确广播。

你的代码里，mask[..., 0]是形状为(480, 360)的2D布尔张量，而你直接用它作为第一个索引去访问形状为(480, 360, 4, 80)的tensor，同时后面又指定了i=2（对应第3个维度）和j=0（对应第4个维度）。这会让PyTorch误以为你想用这个2D掩码去索引前两个维度，但后面又单独指定了第2个维度的索引，导致维度匹配冲突，所以抛出错误。

而NumPy会自动处理这种“部分维度掩码+固定索引”的场景，规则更灵活，所以你的代码在NumPy里能正常运行。

两种可行的解决思路 #

思路1：先固定维度，再应用布尔掩码 #

最简单的方式是先把需要固定的维度（i和j对应的维度）提取出来，得到和掩码形状一致的张量，再用布尔掩码赋值：

i = 2 
j = 0 
mask = torch.randn(480, 360, 3) > 0 
tensor = torch.zeros(480, 360, 4, 80) 

# 先提取dim2=i、dim3=j的切片，得到形状(480,360)的张量
tensor_slice = tensor[..., i, j]
# 用掩码赋值
tensor_slice[mask[..., 0]] = 1
# 或者直接写成一行：
tensor[..., i, j][mask[..., 0]] = 1

这种方式先将张量降到和掩码相同的维度，PyTorch就能正确识别每个掩码位置对应的元素。

思路2：将布尔掩码转换为坐标索引 #

如果你需要更灵活的索引方式，可以用torch.where获取掩码对应的坐标，再用坐标索引赋值：

i = 2 
j = 0 
mask = torch.randn(480, 360, 3) > 0 
tensor = torch.zeros(480, 360, 4, 80) 

# 获取mask[...,0]中为True的位置的行和列坐标
rows, cols = torch.where(mask[..., 0])
# 用坐标索引赋值
tensor[rows, cols, i, j] = 1

这种方式通过显式的坐标索引，完全避免了布尔掩码和固定索引的冲突，适合复杂的索引场景。

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