无法运行 PyTorch 模型（IndexError：维度超出范围（预计在 [-1, 0] 范围内，但得到 1））

Question

我创建了一个 PyTorch 模型并对其进行了训练。这是一个非常简单的神经网络：

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):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

但是，在训练之后，当我尝试使用以下代码运行它时：

model.eval()

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

其中

corpus_train

代表用于训练的

DataLoader

，我不断收到以下错误：

IndexError                                Traceback (most recent call last)
<ipython-input-46-73e4e7e0a291> in <cell line: 4>()
      3 with torch.no_grad():
----> 4   print(model(corpus_test.dataset[0][0]))

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

模型的输入应该是具有 16 个元素的一维张量，然后输出将是 1728 个元素长的张量。

提前致谢！

Answer 1

你尝试过吗：

    print(model(next(iter(corpus_train.dataset))[0]))

参见 https://pytorch.org/tutorials/beginner/basics/data_tutorial.html#iterate-through-the-dataloader