PyTorch基于MNIST训练MLP触发维度越界IndexError的问题求助

核心报错触发原因 #

报错栈定位到nn.Flatten的forward方法抛出IndexError，根源是MLP类初始化时给nn.Flatten传了完全不符合定义的参数：

• nn.Flatten的官方入参仅支持两个维度参数：start_dim（默认为1）、end_dim（默认为-1），不存在名为BATCH的关键字入参。你传入的BATCH=64会被PyTorch解析为flatten操作的起始维度值，对应维度索引64，超出当前4维输入（MNIST图像张量维度范围为[-4, 3]）的可接受索引范围，直接触发维度越界报错。
• 你传入的第二个位置参数input_size也会被解析为end_dim的参数值，进一步放大维度索引不匹配的问题。

关联代码隐患 #

除了核心的Flatten参数错误，现有代码还存在几个会导致运行失败的问题：

• 训练循环里的images = images.permute(0, 3, 1, 2)属于冗余操作。如果使用PyTorch官方MNIST数据集加载器，输出的图像张量本身就是[batch, channel, height, width]的标准NCHW格式，额外permute反而会打乱维度顺序；只有自定义加载器输出NHWC格式时才需要保留这行。
• 代码中Tensorboard的SummaryWriter初始化被注释，但训练循环内直接调用了writer.add_scalar方法，跑完第一个epoch的训练步骤后会触发NameError，需要取消writer初始化的注释，或暂时注释掉日志记录的代码段。

修正后的可运行代码片段 #

修正版MLP结构

class MLP(nn.Module):
    def __init__(self, device, input_size=1*28*28, output_size=10):
        super().__init__()
        
        self.seq = nn.Sequential(
            # 用Flatten默认参数即可：从第1维（通道维）开始展平，到最后一维结束
            # 输入形状[64, 1, 28, 28]经过Flatten后会变成[64, 784]，完全匹配Linear层的输入尺寸要求
            nn.Flatten(),
            nn.Linear(input_size, 32),
            nn.ReLU(),
            nn.Linear(32, output_size)
        )
        self.to(device)
        
    def forward(self, x):
        return self.seq(x)

修正版训练循环片段

for data in train_loader:
    images, labels = data
    images, labels = images.to(device), labels.to(device)
    # 若使用官方MNIST加载器，删除冗余的permute操作
    # images = images.permute(0, 3, 1, 2)
    
    optimizer.zero_grad()
    outputs = model(images)
    loss = criterion(outputs, labels)
    # 后续反向传播、精度计算、验证流程逻辑保持不变

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