I’m getting the error ValueError: Expected input batch_size (1) to match target batch_size (10000). for the following code:
class LogLinearLM(nn.Module):
    def __init__(self, vocab_size):
        super(LogLinearLM, self).__init__()
        self.linear = nn.Linear(2*vocab_size, vocab_size)
        self.softmax = nn.LogSoftmax(dim=1)
    
    def forward(self, input):
        out = self.linear(input)
        return self.softmax(out)


losses = []
loss_fn = nn.NLLLoss()
model = LogLinearLM(10000)
optimizer = optim.SGD(model.parameters(), lr=0.001)
for _ in range(30):
    total_loss = 0
    for w1, w2, w3 in nltk.trigrams(train):
        model.zero_grad()
        X = torch.cat([one_hot(w1), one_hot(w2)])
        yhat = model(X.view(1, -1))
        y = torch.tensor(one_hot(w3), dtype=torch.long)
        loss = loss_fn(yhat, y)
        loss.backward()
        optimizer.step()
        total_loss += loss
    losses.append(total_loss)
plt.plot(losses)

The function one_hot is:
def one_hot(word):
    v = torch.zeros(10000)
    v[word2idx[word]] = 1
    return v

Can anyone help me, please? Thanks