Hi all, I am writing a simple neural network using LSTM to get some understanding of NER. I understand the whole idea but got into trouble with some dimension issues, here’s the problem:

```
class NERModel(nn.Module):
"""
Encoder for NER model.
Args:
- vocab_size: vocabulary size, integer.
- embedding_size: embedding size, integer.
- enc_units: hidden size of LSTM layer, integer.
- ffc_units: hidden units of feedforward layer, integer.
- num_labels: number of named entities. The value should be (actual_num_labels + 1),
because zero paddings are added to the sequences.
"""
def __init__(self, vocab_size, embedding_size, enc_units, ffc_units, num_labels):
super(NERModel, self).__init__()
# Word embedding layer.
self.embedding = nn.Embedding(vocab_size, embedding_size)
# LSTM layer with units of enc_units
self.LSTM = nn.LSTM(embedding_size, enc_units, batch_first=True)
self.dense1 = nn.Linear(enc_units, ffc_units)
self.dense2 = nn.Linear(ffc_units, num_labels)
def forward(self, x):
"""
Args:
- x: Input tensor of shape (batch_size, sequence_length)
Return:
Tensor of shape (batch_size, sequence_length, num_labels)
"""
x = self.embedding(x)
# after embedding: torch.Size([64, 124, 256])
x, _ = self.LSTM(x)
# after lstm: torch.Size([64, 124, 256])
x = self.dense1(x)
# after linear 1: torch.Size([64, 124, 256])
x = self.dense2(x)
# after linear 2: torch.Size([64, 124, 6])
output = F.log_softmax(x, dim=1)
# after softmax: torch.Size([64, 124, 6])
return output
# initialize model
model = NERModel(vocab_size, embedding_size, enc_units, ffc_units, num_labels)
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)
criterion = nn.CrossEntropyLoss()
for i, (value, label) in enumerate(train_loader):
print(value.shape)
optimizer.zero_grad()
outputs = model(value)
# outputs shape: torch.Size([64, 124, 6])
# label shape: torch.Size([64, 124])
loss = criterion(outputs, label)
```

Things all looked good but I got the following error reported:

```
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-9-92530e221aaf> in <module>()
---> 14 loss = criterion(outputs, label)
~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
491 result = self._slow_forward(*input, **kwargs)
492 else:
--> 493 result = self.forward(*input, **kwargs)
494 for hook in self._forward_hooks.values():
495 hook_result = hook(self, input, result)
~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
940 def forward(self, input, target):
941 return F.cross_entropy(input, target, weight=self.weight,
--> 942 ignore_index=self.ignore_index, reduction=self.reduction)
943
944
~/anaconda3/lib/python3.6/site-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
2054 if size_average is not None or reduce is not None:
2055 reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 2056 return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
2057
2058
~/anaconda3/lib/python3.6/site-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
1879 if target.size()[1:] != input.size()[2:]:
1880 raise ValueError('Expected target size {}, got {}'.format(
-> 1881 out_size, target.size()))
1882 input = input.contiguous().view(n, c, 1, -1)
1883 target = target.contiguous().view(n, 1, -1)
ValueError: Expected target size (64, 6), got torch.Size([64, 124])
```

I have outputs shape torch.Size([64, 124, 6]) and label shape: torch.Size([64, 124]). It seems that loss function want me to have outputs shape torch.Size([64, 6, 124]). I don’t get the reason why it is like this, and can some one tell me how to modify it?