在fastNLP框架中加入LSTM遇到的一个问题

Traceback (most recent call last):
File “train_tener_cn.py”, line 176, in
use_tqdm=True, print_every=300, save_path=None)
File “/usr/local/lib/python3.6/dist-packages/fastNLP/core/trainer.py”, line 520, in init
batch_size=check_batch_size)
File “/usr/local/lib/python3.6/dist-packages/fastNLP/core/trainer.py”, line 920, in _check_code
pred_dict = model(**refined_batch_x)
File “/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py”, line 532, in call
result = self.forward(*input, **kwargs)
File “/content/drive/My Drive/NER/TENER-weibo_msra/models/TENER.py”, line 105, in forward
return self._forward(words, target, bigrams)
File “/content/drive/My Drive/NER/TENER-weibo_msra/models/TENER.py”, line 86, in forward
words,(hidden,) = self.lstm(words,seq_len=words.size(1))
File “/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py”, line 532, in call
result = self.forward(*input, **kwargs)
File “/usr/local/lib/python3.6/dist-packages/fastNLP/modules/encoder/lstm.py”, line 68, in forward
sort_lens, sort_idx = torch.sort(seq_len, dim=0, descending=True)
TypeError: sort() received an invalid combination of arguments - got (int, descending=bool, dim=int), but expected one of:

• (Tensor input, name dim, bool descending, tuple of Tensors out)
• (Tensor input, int dim, bool descending, tuple of Tensors out)

words,hidden = self.lstm(words)

问题原因：
LSTM中的输入words对应矩阵----[batch_size, seq_len, feature_dim]
其中，batch_size=16,feature_dim=768, seq_len是句子长度，在不断变化。
但是fastnlp框架下，一开始会随机选择两个例子来进行测试，防止出错，如矩阵为[2,26]，这样的话就与我的设置不符，导致了出错。

解决办法：
加入一个判断条件，如果batch_size==16,则加入LSTM，否则不加入。