别再被 PyTorch 维度报错折磨了，其实关键就在 keepdim 这个参数上
RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at dimension 1。很多初学者在处理归一化或计算均值时，习惯性地写 x.sum(dim=1)，结果要么直接程序崩溃，要么更糟糕——代码居然跑通了，但算出来的数值完全不对，导致模型训练不出结果。其实这背后隐藏着 PyTorch 维度压缩与广播机制（Broadcasting）的深层逻辑。
在 PyTorch 中，所有的还原操作（Reduction Operations），比如 sum()、mean()、max() 或 min()，其默认行为都是将指定的维度“压扁”并直接删除。举个具体的例子：假设你有一个形状为 (2, 3) 的张量，你在 dim=1 上执行求和操作。默认情况下，PyTorch 会把这个维度直接删掉，最终输出的形状会变成 (2,)。
如果你在操作时设置了 keepdim=True，这个维度虽然数值上被压缩了，但其位置会被保留，形状会变成 (2, 1)。虽然在数值上两者看起来一样，但在张量运算中，这一个维度的差异就是“正确结果”与“逻辑 Bug”的分水岭。
我们可以通过一段简单的实操对比来看清这个差异：
import torch
# 创建一个 2x3 的张量
m = torch.tensor([[1., 2., 3.],
[4., 5., 6.]])
# 默认情况：维度被直接删除
res_default = m.sum(dim=1)
print(res_default.shape) # 输出: torch.Size([2])
# 开启 keepdim：维度被保留为 1
res_keep = m.sum(dim=1, keepdim=True)
print(res_keep.shape) # 输出: torch.Size([2, 1])真正决定成败的是接下来的“广播机制”。当你需要用原张量除以这个求和结果（例如进行行归一化）时，PyTorch 会尝试从右向左对齐两个张量的维度。
如果使用了 keepdim=True，原张量 (2, 3) 与结果张量 (2, 1) 对齐，PyTorch 发现最后一个维度是 3 和 1，此时广播机制会自动将 1 拉伸为 3，从而实现完美的逐元素计算。
但如果你没加这个参数，原张量 (2, 3) 将面对一个形状为 (2,) 的张量。此时 PyTorch 对齐维度时发现末尾是 3 vs 2，由于不符合广播规则，直接抛出 RuntimeError 导致程序崩溃。
# 正确做法：keepdim=True 保证能够正确广播
normed = m / m.sum(dim=1, keepdim=True)
print(normed.sum(dim=1)) # 输出: tensor([1., 1.]) ✅ 每一行和正确归一化为 1这里有一个最隐蔽的“坑”：如果你的张量刚好是方阵（比如 (2, 2)），不加 keepdim 竟然可能不会报错。因为 (2, 2) 与 (2,) 在广播时，PyTorch 可能会将其误认为是在另一个维度上进行操作，导致计算逻辑完全错误，但程序却能顺利跑完。这种没有报错的数值 Bug 是最难调试的，因为它不会在控制台留下任何线索。
总结起来，只要你的业务场景涉及到“先对某个维度降维求值，随后又要将该结果与原张量进行算术运算”，请务必养成闭眼加上 keepdim=True 的习惯。这不仅能避免低级的维度报错，更能有效杜绝那些潜伏在代码深处的数值逻辑错误。
