class LSTMTagger(nn.Module):
def __init__(self, embedding_dim, hidden_dim, vocab_size, tagset_size):
    super(LSTMTagger, self).__init__()
    self.hidden_dim = hidden_dim

    self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)

    # The LSTM takes word embeddings as inputs, and outputs hidden states
    # with dimensionality hidden_dim.
    self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first = True)

    # The linear layer that maps from hidden state space to tag space
    self.hidden2tag = nn.Linear(hidden_dim, tagset_size)

def forward(self, sentence):
    batch_size = sentence.size(0)
    sentence = sentence.long()
    embeds = self.word_embeddings(sentence)
    lstm_out, _ = self.lstm(embeds)
    lstm_out = lstm_out.contiguous().view(-1, self.hidden_dim)
    tag_space = self.hidden2tag(lstm_out)
    tag_scores = F.log_softmax(tag_space, dim=1)
    return tag_scores

EMBEDDING_DIM = 32
HIDDEN_DIM = 32
EPOCHS = 10
#len(word_to_ix) = vocab_size, len(tag_to_ix) = tagset_size
model = LSTMTagger(EMBEDDING_DIM, HIDDEN_DIM, len(word_to_ix), len(tag_to_ix))
loss_function = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
for epoch in range(EPOCHS):
for sentence, tags in train_loader:
model.zero_grad()
    tag_scores = model(sentence)
    loss = loss_function(tag_scores, tags.flatten())
    loss.backward()
    optimizer.step()

The above code works now. Please look into it and tell me is there any fault?