My criterion is down below. As for the forward function, I didn’t copy the entire function. I edit the post just now
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(lstm_model.parameters(), lr=learning_rate)