# Buggy code
def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
    """Add embeddings to all words in a list of sentences. If embeddings are already added,
    updates only if embeddings are non-static."""

    # first, find longest sentence in batch
    longest_sentence_in_batch: int = len(
        max(
            [
                self.tokenizer.tokenize(sentence.to_tokenized_string())
                for sentence in sentences
            ],
            key=len,
        )
    )

    # prepare id maps for BERT model
    features = self._convert_sentences_to_features(sentences, longest_sentence_in_batch)
    if torch.cuda.is_available():
        all_input_ids = torch.tensor(
            [f.input_ids for f in features], dtype=torch.long
        ).cuda()
        all_input_masks = torch.tensor(
            [f.input_mask for f in features], dtype=torch.long
        ).cuda()
    else:
        all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
        all_input_masks = torch.tensor(
            [f.input_mask for f in features], dtype=torch.long
        )

    # put encoded batch through BERT model to get all hidden states of all encoder layers
    self.model.eval()
    all_encoder_layers, _ = self.model(
        all_input_ids, token_type_ids=None, attention_mask=all_input_masks
    )

    with torch.no_grad():
        for sentence_index, sentence in enumerate(sentences):
            feature = features[sentence_index]

            # get aggregated embeddings for each BERT-subtoken in sentence
            subtoken_embeddings = []
            for token_index, _ in enumerate(feature.tokens):
                all_layers = []
                for layer_index in self.layer_indexes:
                    layer_output = (
                        all_encoder_layers[int(layer_index)]
                        .detach()
                        .cpu()[sentence_index]
                    )
                    all_layers.append(layer_output[token_index])

                subtoken_embeddings.append(torch.cat(all_layers))

            # get the current sentence object
            token_idx = 0
            for token in sentence:
                # add concatenated embedding to sentence
                token_idx += 1

                if self.pooling_operation == "first":
                    # use first subword embedding if pooling operation is 'first'
                    token.set_embedding(self.name, subtoken_embeddings[token_idx])
                else:
                    # otherwise, do a mean over all subwords in token
                    embeddings = subtoken_embeddings[
                        token_idx : token_idx + feature.token_subtoken_count[token.idx]
                    ]
                    embeddings = [embedding.unsqueeze(0) for embedding in embeddings]
                    mean = torch.mean(torch.cat(embeddings, dim=0), dim=0)
                    token.set_embedding(self.name, mean)

                token_idx += feature.token_subtoken_count[token.idx] - 1

    return sentences


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-1-47fbfcb3fd4f> in <module>
5 BertEmbedding = BertEmbeddings()
6 sent = Sentence('I love Austin')
----> 7 BertEmbedding.embed(sent)
8 # fe.embed(sent)

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/flair/embeddings.py in embed(self, sentences)
54
55         if not everything_embedded or not self.static_embeddings:
---> 56             self._add_embeddings_internal(sentences)
57
58         return sentences

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/flair/embeddings.py in _add_embeddings_internal(self, sentences)
957         # put encoded batch through BERT model to get all hidden states of all encoder layers
958         self.model.eval()
--> 959         all_encoder_layers, _ = self.model(all_input_ids, token_type_ids=None, attention_mask=all_input_masks)
960
961         with torch.no_grad():

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
487             result = self._slow_forward(*input, **kwargs)
488         else:
--> 489             result = self.forward(*input, **kwargs)
490         for hook in self._forward_hooks.values():
491             hook_result = hook(self, input, result)

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/pytorch_pretrained_bert/modeling.py in forward(self, input_ids, token_type_ids, attention_mask, output_all_encoded_layers)
607         extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
608
--> 609         embedding_output = self.embeddings(input_ids, token_type_ids)
610         encoded_layers = self.encoder(embedding_output,
611                                       extended_attention_mask,

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
487             result = self._slow_forward(*input, **kwargs)
488         else:
--> 489             result = self.forward(*input, **kwargs)
490         for hook in self._forward_hooks.values():
491             hook_result = hook(self, input, result)

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/pytorch_pretrained_bert/modeling.py in forward(self, input_ids, token_type_ids)
191             token_type_ids = torch.zeros_like(input_ids)
192
--> 193         words_embeddings = self.word_embeddings(input_ids)
194         position_embeddings = self.position_embeddings(position_ids)
195         token_type_embeddings = self.token_type_embeddings(token_type_ids)

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
487             result = self._slow_forward(*input, **kwargs)
488         else:
--> 489             result = self.forward(*input, **kwargs)
490         for hook in self._forward_hooks.values():
491             hook_result = hook(self, input, result)

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/torch/nn/modules/sparse.py in forward(self, input)
116         return F.embedding(
117             input, self.weight, self.padding_idx, self.max_norm,
--> 118             self.norm_type, self.scale_grad_by_freq, self.sparse)
119
120     def extra_repr(self):

~/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/torch/nn/functional.py in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
1452         # remove once script supports set_grad_enabled
1453         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
-> 1454     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
1455
1456

RuntimeError: Expected object of backend CPU but got backend CUDA for argument #3 'index'

# Fixed code
def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
    """Add embeddings to all words in a list of sentences. If embeddings are already added,
    updates only if embeddings are non-static."""

    # first, find longest sentence in batch
    longest_sentence_in_batch: int = len(
        max(
            [
                self.tokenizer.tokenize(sentence.to_tokenized_string())
                for sentence in sentences
            ],
            key=len,
        )
    )

    # prepare id maps for BERT model
    features = self._convert_sentences_to_features(sentences, longest_sentence_in_batch)
    if torch.cuda.is_available():
        all_input_ids = torch.tensor(
            [f.input_ids for f in features], dtype=torch.long
        ).cuda()
        all_input_masks = torch.tensor(
            [f.input_mask for f in features], dtype=torch.long
        ).cuda()
    else:
        all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
        all_input_masks = torch.tensor(
            [f.input_mask for f in features], dtype=torch.long
        )

    # put encoded batch through BERT model to get all hidden states of all encoder layers
    if torch.cuda.is_available():
        self.model.cuda()
    self.model.eval()
    all_encoder_layers, _ = self.model(
        all_input_ids, token_type_ids=None, attention_mask=all_input_masks
    )

    with torch.no_grad():
        for sentence_index, sentence in enumerate(sentences):
            feature = features[sentence_index]

            # get aggregated embeddings for each BERT-subtoken in sentence
            subtoken_embeddings = []
            for token_index, _ in enumerate(feature.tokens):
                all_layers = []
                for layer_index in self.layer_indexes:
                    layer_output = (
                        all_encoder_layers[int(layer_index)]
                        .detach()
                        .cpu()[sentence_index]
                    )
                    all_layers.append(layer_output[token_index])

                subtoken_embeddings.append(torch.cat(all_layers))

            # get the current sentence object
            token_idx = 0
            for token in sentence:
                # add concatenated embedding to sentence
                token_idx += 1

                if self.pooling_operation == "first":
                    # use first subword embedding if pooling operation is 'first'
                    token.set_embedding(self.name, subtoken_embeddings[token_idx])
                else:
                    # otherwise, do a mean over all subwords in token
                    embeddings = subtoken_embeddings[
                        token_idx : token_idx + feature.token_subtoken_count[token.idx]
                    ]
                    embeddings = [embedding.unsqueeze(0) for embedding in embeddings]
                    mean = torch.mean(torch.cat(embeddings, dim=0), dim=0)
                    token.set_embedding(self.name, mean)

                token_idx += feature.token_subtoken_count[token.idx] - 1

    return sentences