PyTorch模型运行报错：IndexError维度超出范围问题求助

问题根源 #

报错IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)的核心原因是输入张量维度与模型线性层要求不匹配：

• 模型第一个线性层nn.Linear(NUM_RANDOMS, 128)要求输入最后一个维度为NUM_RANDOMS=16，且输入应为2D张量（形状为[batch_size, 16]）。
• 直接从dataset取出的单个样本corpus_test.dataset[0][0]是1D张量（形状为[16]），PyTorch会误将其识别为[16, 1]（即16个样本、每个样本1个特征），与模型期望的输入维度冲突。

解决方法 #

方法1：手动添加batch维度 #

在喂给模型前，用unsqueeze(0)给输入张量增加一个batch维度，将形状从[16]转为[1, 16]：

model.eval()

with torch.no_grad():
    input_tensor = corpus_test.dataset[0][0].unsqueeze(0)
    print(model(input_tensor))

方法2：修改模型兼容1D输入 #

如果希望模型同时支持单个1D样本和批量2D样本，可在forward方法中自动处理维度：

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(NUM_RANDOMS, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, MAX_NUM_OUTPUT),
        )

    def forward(self, x):
        # 若输入为1D张量，自动添加batch维度
        if x.dim() == 1:
            x = x.unsqueeze(0)
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

修改后直接传入1D张量即可正常运行：

model.eval()

with torch.no_grad():
    print(model(corpus_test.dataset[0][0]))

方法3：通过DataLoader获取输入（推荐） #

训练时用DataLoader自动打包batch，推理时也可以沿用该方式，从根源避免维度问题：

model.eval()

with torch.no_grad():
    for batch in corpus_test:
        # batch[0]是批量的random_tensor，形状为[1, 16]（默认batch_size=1）
        print(model(batch[0]))
        break  # 仅查看第一个样本

验证输入形状 #

可以先打印张量形状确认问题：

print(corpus_test.dataset[0][0].shape)  # 输出: torch.Size([16])
print(corpus_test.dataset[0][0].unsqueeze(0).shape)  # 输出: torch.Size([1, 16])

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