M2 Mac上PyTorch用MPS抛出“张量类型不匹配”错误的原因？

核心问题排查 #

你遇到的"Mismatched Tensor types in NNPack convolutionOutput"错误，主要由以下操作导致：

1. Softmax维度错误 + CrossEntropyLoss的冗余使用 #

模型最后一层的nn.Softmax(dim=0)存在两个问题：

• CrossEntropyLoss内部已整合LogSoftmax与NLLLoss，额外添加Softmax会导致输出分布与标签维度不匹配，MPS设备的严格类型校验会触发报错。
• dim=0是按样本维度做归一化，而MNIST任务应按类别维度（dim=1）计算，错误的归一化逻辑会进一步加剧张量类型/维度不匹配。

2. Test循环中张量设备不一致 #

test_loop函数未将测试数据X和y移动到MPS设备，而模型已部署在MPS上，跨设备的张量运算会直接触发类型不匹配错误。

3. MPS兼容性细节 #

PyTorch早期版本对MPS的算子支持存在严格的类型校验要求，比如卷积层输入与权重必须为同类型张量（均为float32），隐式类型转换（如Dropout在MPS上的行为）可能触发报错。

修复步骤 #

步骤1：移除冗余的Softmax层 #

修改NeuralNetwork类：

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.mnist_nn = nn.Sequential(
            nn.Conv2d(1, 32, 3),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3),
            nn.ReLU(),
            nn.MaxPool2d(2,stride=2),
            nn.Dropout(p=0.25),
            nn.Flatten(),
            nn.Linear(9216, 128),
            nn.ReLU(),
            nn.Dropout(p=0.25),
            nn.Linear(128, 10)  # 移除Softmax层
        )

    def forward(self, x):
        logits = self.mnist_nn(x)
        return logits

步骤2：统一Test循环的设备 #

更新test_loop函数，将数据同步到MPS：

def test_loop(dataloader, model, loss_fn):
    model.eval()
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            X = X.to(DEVICE)  # 新增：将输入移到MPS
            y = y.to(DEVICE)  # 新增：将标签移到MPS
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"测试误差：\n 准确率: {(100*correct):>0.1f}%, 平均损失: {test_loss:>8f} \n")

步骤3：验证PyTorch版本 #

确保使用PyTorch 2.0及以上版本，MPS兼容性在该版本后有大幅提升，检查版本命令：

python -c "import torch; print(torch.__version__)"

额外建议 #

若修改后仍报错，可显式指定张量类型为float32，比如在数据加载后添加X = X.to(DEVICE, dtype=torch.float32)，避免隐式类型转换引发问题。

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