I want to write a sequence tagger, or NER you could say, with an (Bi-)LSTM.
This is my code (provide the whole just fyi, but the most important part is the LSTMTagger class on the top and the training loop at the bottom:
import sys, os
from torchnlp.word_to_vector import FastText
from pathlib import Path
import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.optim as optim
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence, pad_sequence
from torch.utils.data import DataLoader, random_split
import torch.nn.functional as F
import re
from torchtext import data
from tqdm import tqdm, trange
import fasttext
import fasttext.util

EMBEDDING_DIM = 300
HIDDEN_DIM = 6
learning_rate = 0.01
epochs = 5
batch_size = 1

tag_to_ix = {"O": 0, "B-SKILL": 1, "I-SKILL": 2}

cuda_device = 0
device = torch.device("cuda:%d" % cuda_device if torch.cuda.is_available() else "cpu")

base_path = Path("..")
fasttext_model = base_path / "models" / "cc.de.300.bin"

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)

        self.lstm_forward = nn.LSTM(embedding_dim, hidden_dim,
                                   batch_first=True,
                                   bidirectional=False)
        
        self.lstm_backward = nn.LSTM(embedding_dim, hidden_dim,
                                    batch_first=True,
                                    bidirectional=False)

        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)

    def forward(self, embeds_forward, embeds_backward):

        lstm_out_forward, _ = self.lstm_forward(embeds_forward)
        # lstm_out_backward, _ = self.lstm_backward(embeds_backward)

        # lstm_concat = torch.cat((lstm_out_forward, lstm_out_backward), 2)
        
        tag_space = self.hidden2tag(lstm_out_forward)
        tag_scores = F.log_softmax(tag_space, dim=1)
        print(tag_scores.size()) # >> torch.Size([1, 2, 3])
        return tag_scores

class EmbeddingVectorizer:
    def __init__(self):

        self.embedding_model = fasttext.load_model(str(fasttext_model))

    def __call__(self, doc):
        """
        Convert address to embedding vectors
        :param address: The address to convert
        :return: The embeddings vectors
        """
        embeddings = []
        for word in doc:
            embeddings.append(self.embedding_model[word])
        return embeddings

embedding_model = EmbeddingVectorizer()

def pad_collate_fn(batch):
    sequences_vectors, sequences_labels, lengths = zip(*[
        (torch.FloatTensor(seq_vectors), torch.LongTensor(labels), len(seq_vectors))
        for (seq_vectors, labels) in sorted(batch, key=lambda x: len(x[0]), reverse=True)
    ])

    lengths = torch.LongTensor(lengths)

    padded_sequences_vectors = pad_sequence(sequences_vectors, batch_first=True, padding_value=0)

    padded_sequences_labels = pad_sequence(sequences_labels, batch_first=True, padding_value=-100)

    return (padded_sequences_vectors, lengths), padded_sequences_labels

def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)

def read_conll(file_path):
    file_path = Path(file_path)

    raw_text = file_path.read_text().strip()
    raw_docs = re.split(r'\n\t?\n', raw_text)
    token_docs = []
    tag_docs = []
    for doc in raw_docs:
        tokens = []
        tags = []
        for line in doc.split('\n'):
            token, tag = line.split('\t')
            tokens.append(token)
            tags.append(tag)
        token_docs.append(tokens)
        tag_docs.append(tags)

    return token_docs, tag_docs

start_path = Path("..")
data_path = start_path / "data" / "traindata_toy.conll"
texts, tags = read_conll(data_path)

training_data = []
for doc, doc_tags in zip(texts, tags):
    training_data.append((doc, doc_tags))  # Example data: [(["a", "b", "c"], ["O", "O", "B-SKILL"]), ([...], [...])]

word_to_ix = {}
for sent, tags in training_data:
    for word in sent:
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)
    
embeddings_tags = []    

for doc, doc_tags in training_data:
    embs = embedding_model(doc)
    targets = prepare_sequence(doc_tags, tag_to_ix)
    embeddings_tags.append((embs, targets))

train_size = int(0.8 * len(embeddings_tags))
valid_size = int((len(embeddings_tags) - train_size) / 2)
test_size = valid_size
train_dataset, valid_dataset, test_dataset = torch.utils.data.random_split(embeddings_tags, [train_size, valid_size, test_size])

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=pad_collate_fn)
valid_loader = DataLoader(valid_dataset, batch_size=batch_size, collate_fn=pad_collate_fn)
test_loader = DataLoader(test_dataset, batch_size=batch_size, collate_fn=pad_collate_fn)


model = LSTMTagger(EMBEDDING_DIM, HIDDEN_DIM, len(word_to_ix), len(tag_to_ix))
loss_function = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)


for epoch in trange(epochs, desc="Epoch"):
    for step, batch in enumerate(train_loader):

        docs = batch[0][0].to(device)
        tags = batch[1].to(device)
        
        model.zero_grad()
        
        doc_forward = docs
        doc_backward = torch.clone(doc_forward)
        doc_backward = torch.flip(doc_backward, [1])

        print(doc_forward.size())  # >> torch.Size([1, 2, 300])

        tag_scores = model(doc_forward, doc_backward)

        print(tags.size()) # >> torch.Size([1, 2])
        print(tag_scores.size()) # >> torch.Size([1, 2, 3])
        loss = loss_function(tag_scores, tags)
        loss.backward()
        optimizer.step()

I get the error
RuntimeError: Expected target size [1, 3], got [1, 2]
at
loss = loss_function(tag_scores, tags).
But the sizes should be right for NLLLoss or not? You can see the sizes I got for a document with two tokens. So the size of the target tags is [1, 2] for one batch and two tokens, and the size of the predicted tags after log_softmax is [1, 2, 3] for three possible tags, as I understand it correctly. But I still get this error. Do you have any idea? Thank you very much!