如何在PyTorch Sequantial模型中指定批次大小？

我在PyTorch中有一个顺序模型：

model = nn.Sequential(
        nn.Embedding(alphabet_size, 64),
        nn.LSTM(64, ...),
        nn.Flatten(),
        nn.Linear(...),
        nn.Softmax()
    )

我想强制将批量大小设置为嵌入层：

nn.Embedding(alphabet_size, 64)

与Keras一样：

Embedding(alphabet_size, 64, batch_input_shape=(batch_size, time_steps)))

---

怎么做？

这是你问题的答案。 下面已经提供了示例