# Buggy code
def __init__(
    self,
    embeddings: List[TokenEmbeddings],
    hidden_size=128,
    rnn_layers=1,
    reproject_words: bool = True,
    reproject_words_dimension: int = None,
    bidirectional: bool = False,
    dropout: float = 0.5,
    word_dropout: float = 0.0,
    locked_dropout: float = 0.0,
    rnn_type="GRU",
):
    """The constructor takes a list of embeddings to be combined.
    :param embeddings: a list of token embeddings
    :param hidden_size: the number of hidden states in the rnn
    :param rnn_layers: the number of layers for the rnn
    :param reproject_words: boolean value, indicating whether to reproject the token embeddings in a separate linear
    layer before putting them into the rnn or not
    :param reproject_words_dimension: output dimension of reprojecting token embeddings. If None the same output
    dimension as before will be taken.
    :param bidirectional: boolean value, indicating whether to use a bidirectional rnn or not
    :param dropout: the dropout value to be used
    :param word_dropout: the word dropout value to be used, if 0.0 word dropout is not used
    :param locked_dropout: the locked dropout value to be used, if 0.0 locked dropout is not used
    :param rnn_type: 'GRU', 'LSTM',  'RNN_TANH' or 'RNN_RELU'
    """
    super().__init__()

    self.embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embeddings)

    self.rnn_type = rnn_type

    self.reproject_words = reproject_words
    self.bidirectional = bidirectional

    self.length_of_all_token_embeddings: int = self.embeddings.embedding_length

    self.static_embeddings = False

    self.__embedding_length: int = hidden_size
    if self.bidirectional:
        self.__embedding_length *= 4

    self.embeddings_dimension: int = self.length_of_all_token_embeddings
    if self.reproject_words and reproject_words_dimension is not None:
        self.embeddings_dimension = reproject_words_dimension

    # bidirectional RNN on top of embedding layer
    self.word_reprojection_map = torch.nn.Linear(
        self.length_of_all_token_embeddings, self.embeddings_dimension
    )
    self.rnn = torch.nn.RNNBase(
        rnn_type,
        self.embeddings_dimension,
        hidden_size,
        num_layers=rnn_layers,
        bidirectional=self.bidirectional,
    )

    self.name = "document_" + self.rnn._get_name()

    # dropouts
    if locked_dropout > 0.0:
        self.dropout: torch.nn.Module = LockedDropout(locked_dropout)
    else:
        self.dropout = torch.nn.Dropout(dropout)

    self.use_word_dropout: bool = word_dropout > 0.0
    if self.use_word_dropout:
        self.word_dropout = WordDropout(word_dropout)

    torch.nn.init.xavier_uniform_(self.word_reprojection_map.weight)

    self.to(flair.device)


# Traceback (KeyError)
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-19-e27894ad7da5> in <module>
1 assert torch.cuda.is_available(), "Do not run this model with GPU support."
2 # Run model
----> 3 model(config)

~/cloud/model.py in model(config, language, optimise)
77                   anneal_factor=config.anneal_factor,
78                   patience=config.patience,
---> 79                   max_epochs=config.max_epochs)
80
81

~/anaconda3/envs/py36/lib/python3.6/site-packages/flair/trainers/trainer.py in train(self, base_path, evaluation_metric, learning_rate, mini_batch_size, eval_mini_batch_size, max_epochs, anneal_factor, patience, train_with_dev, monitor_train, embeddings_in_memory, checkpoint, save_final_model, anneal_with_restarts, shuffle, param_selection_mode, num_workers, **kwargs)
91                 [Sentence("d", labels=["0.1"])],
92                 eval_mini_batch_size,
---> 93                 embeddings_in_memory,
94             )
95             if log_train:

~/anaconda3/envs/py36/lib/python3.6/site-packages/flair/models/text_classification_model.py in evaluate(self, sentences, eval_mini_batch_size, embeddings_in_memory, out_path)
175                 batch_count += 1
176
--> 177                 labels, loss = self.forward_labels_and_loss(batch)
178
179                 clear_embeddings(

~/anaconda3/envs/py36/lib/python3.6/site-packages/flair/models/text_classification_model.py in forward_labels_and_loss(self, sentences)
106         self, sentences: Union[Sentence, List[Sentence]]
107     ) -> (List[List[Label]], torch.tensor):
--> 108         scores = self.forward(sentences)
109         labels = self._obtain_labels(scores)
110         loss = self._calculate_loss(scores, sentences)

~/anaconda3/envs/py36/lib/python3.6/site-packages/flair/models/text_classification_model.py in forward(self, sentences)
68
69     def forward(self, sentences) -> List[List[float]]:
---> 70         self.document_embeddings.embed(sentences)
71
72         text_embedding_list = [

~/anaconda3/envs/py36/lib/python3.6/site-packages/flair/embeddings.py in embed(self, sentences)
2079         self.rnn.flatten_parameters()
2080
-> 2081         rnn_out, hidden = self.rnn(packed)
2082
2083         outputs, output_lengths = torch.nn.utils.rnn.pad_packed_sequence(rnn_out)

~/anaconda3/envs/py36/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/envs/py36/lib/python3.6/site-packages/torch/nn/modules/rnn.py in forward(self, input, hx)
206
207         self.check_forward_args(input, hx, batch_sizes)
--> 208         _impl = _rnn_impls[self.mode]
209         if batch_sizes is None:
210             result = _impl(input, hx, self._get_flat_weights(), self.bias, self.num_layers,

KeyError: 'LSTM'

# Fixed code
def __init__(
    self,
    embeddings: List[TokenEmbeddings],
    hidden_size=128,
    rnn_layers=1,
    reproject_words: bool = True,
    reproject_words_dimension: int = None,
    bidirectional: bool = False,
    dropout: float = 0.5,
    word_dropout: float = 0.0,
    locked_dropout: float = 0.0,
    rnn_type="GRU",
):
    """The constructor takes a list of embeddings to be combined.
    :param embeddings: a list of token embeddings
    :param hidden_size: the number of hidden states in the rnn
    :param rnn_layers: the number of layers for the rnn
    :param reproject_words: boolean value, indicating whether to reproject the token embeddings in a separate linear
    layer before putting them into the rnn or not
    :param reproject_words_dimension: output dimension of reprojecting token embeddings. If None the same output
    dimension as before will be taken.
    :param bidirectional: boolean value, indicating whether to use a bidirectional rnn or not
    :param dropout: the dropout value to be used
    :param word_dropout: the word dropout value to be used, if 0.0 word dropout is not used
    :param locked_dropout: the locked dropout value to be used, if 0.0 locked dropout is not used
    :param rnn_type: 'GRU' or 'LSTM'
    """
    super().__init__()

    self.embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embeddings)

    self.rnn_type = rnn_type

    self.reproject_words = reproject_words
    self.bidirectional = bidirectional

    self.length_of_all_token_embeddings: int = self.embeddings.embedding_length

    self.static_embeddings = False

    self.__embedding_length: int = hidden_size
    if self.bidirectional:
        self.__embedding_length *= 4

    self.embeddings_dimension: int = self.length_of_all_token_embeddings
    if self.reproject_words and reproject_words_dimension is not None:
        self.embeddings_dimension = reproject_words_dimension

    self.word_reprojection_map = torch.nn.Linear(
        self.length_of_all_token_embeddings, self.embeddings_dimension
    )

    # bidirectional RNN on top of embedding layer
    if rnn_type == "LSTM":
        self.rnn = torch.nn.LSTM(
            self.embeddings_dimension,
            hidden_size,
            num_layers=rnn_layers,
            bidirectional=self.bidirectional,
        )
    else:
        self.rnn = torch.nn.GRU(
            self.embeddings_dimension,
            hidden_size,
            num_layers=rnn_layers,
            bidirectional=self.bidirectional,
        )

    self.name = "document_" + self.rnn._get_name()

    # dropouts
    if locked_dropout > 0.0:
        self.dropout: torch.nn.Module = LockedDropout(locked_dropout)
    else:
        self.dropout = torch.nn.Dropout(dropout)

    self.use_word_dropout: bool = word_dropout > 0.0
    if self.use_word_dropout:
        self.word_dropout = WordDropout(word_dropout)

    torch.nn.init.xavier_uniform_(self.word_reprojection_map.weight)

    self.to(flair.device)