train.py

losses = []
loss_function = nn.NLLLoss()
model = LeModel(len(vocab))
optimizer = optim.SGD(model.parameters(), lr=0.001)

for epoch in tqdm(range(1), leave=False):
    total_loss = 0
    for context, target in tqdm(data, leave=False):
        model.zero_grad()
        
        predictions = model(x_train)
        loss = loss_function(predictions, y_train, dtype=torch.long)

        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    print("%s: %s" % (epoch, loss.item()))

    losses.append(total_loss)
    torch.save(model.state_dict(), "model.torch")