# Buggy code
def __setstate__(self, state):
    from bpemb.util import sentencepiece_load

    model_file = self.model_tpl.format(lang=state["lang"], vs=state["vs"])
    self.__dict__ = state

    # write out the binary sentence piece model into the expected directory
    self.cache_dir: Path = Path(flair.cache_root) / "embeddings"
    if "spm_model_binary" in self.__dict__:
        # if the model was saved as binary and it is not found on disk, write to appropriate path
        if not os.path.exists(self.cache_dir / state["lang"]):
            os.makedirs(self.cache_dir / state["lang"])
        self.model_file = self.cache_dir / model_file
        with open(self.model_file, "wb") as out:
            out.write(self.__dict__["spm_model_binary"])
    else:
        # otherwise, use normal process and potentially trigger another download
        self.model_file = self._load_file(model_file)

    # once the modes if there, load it with sentence piece
    state["spm"] = sentencepiece_load(self.model_file)


# Traceback (RuntimeError)
RuntimeError                              Traceback (most recent call last)
<ipython-input-21-83d7f894d041> in <module>()

---->  classifier.predict(sentence)

12 frames
/usr/local/lib/python3.6/dist-packages/flair/models/text_classification_model.py in predict(self, sentences, mini_batch_size, multi_class_prob, verbose, label_name, return_loss, embedding_storage_mode)
220                     continue
221
--> 222                 scores = self.forward(batch)
223
224                 if return_loss:

/usr/local/lib/python3.6/dist-packages/flair/models/text_classification_model.py in forward(self, sentences)
97     def forward(self, sentences):
98
---> 99         self.document_embeddings.embed(sentences)
100
101         embedding_names = self.document_embeddings.get_names()

/usr/local/lib/python3.6/dist-packages/flair/embeddings/legacy.py in embed(self, sentences)
1423         sentences.sort(key=lambda x: len(x), reverse=True)
1424
-> 1425         self.embeddings.embed(sentences)
1426
1427         # first, sort sentences by number of tokens

/usr/local/lib/python3.6/dist-packages/flair/embeddings/token.py in embed(self, sentences, static_embeddings)
69
70         for embedding in self.embeddings:
---> 71             embedding.embed(sentences)
72
73     @property

/usr/local/lib/python3.6/dist-packages/flair/embeddings/base.py in embed(self, sentences)
59
60         if not everything_embedded or not self.static_embeddings:
---> 61             self._add_embeddings_internal(sentences)
62
63         return sentences

/usr/local/lib/python3.6/dist-packages/flair/embeddings/token.py in _add_embeddings_internal(self, sentences)
1702             sentence_words.append([token.text for token in sentence])
1703
-> 1704         embeddings = self.ee.embed_batch(sentence_words)
1705
1706         for i, sentence in enumerate(sentences):

/usr/local/lib/python3.6/dist-packages/allennlp/commands/elmo.py in embed_batch(self, batch)
253         if batch == [[]]:
254             elmo_embeddings.append(empty_embedding())
--> 255         else:
256             embeddings, mask = self.batch_to_embeddings(batch)
257             for i in range(len(batch)):

/usr/local/lib/python3.6/dist-packages/allennlp/commands/elmo.py in batch_to_embeddings(self, batch)
195         if self.cuda_device >= 0:
196             character_ids = character_ids.cuda(device=self.cuda_device)
--> 197
198         bilm_output = self.elmo_bilm(character_ids)
199         layer_activations = bilm_output['activations']

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
548             result = self._slow_forward(*input, **kwargs)
549         else:
--> 550             result = self.forward(*input, **kwargs)
551         for hook in self._forward_hooks.values():
552             hook_result = hook(self, input, result)

/usr/local/lib/python3.6/dist-packages/allennlp/modules/elmo.py in forward(self, inputs, word_inputs)
605                 type_representation = token_embedding['token_embedding']
606         else:
--> 607             token_embedding = self._token_embedder(inputs)
608             mask = token_embedding['mask']
609             type_representation = token_embedding['token_embedding']

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
548             result = self._slow_forward(*input, **kwargs)
549         else:
--> 550             result = self.forward(*input, **kwargs)
551         for hook in self._forward_hooks.values():
552             hook_result = hook(self, input, result)

/usr/local/lib/python3.6/dist-packages/allennlp/modules/elmo.py in forward(self, inputs)
357         character_embedding = torch.nn.functional.embedding(
358                 character_ids_with_bos_eos.view(-1, max_chars_per_token),
--> 359                 self._char_embedding_weights
360         )
361

/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
1722         # remove once script supports set_grad_enabled
1723         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
-> 1724     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
1725
1726

RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_index_select

# Fixed code
def __setstate__(self, state):
    self.__dict__ = state

    if re.fullmatch(r"cuda:[0-9]+", str(flair.device)):
        cuda_device = int(str(flair.device).split(":")[-1])
    elif str(flair.device) == "cpu":
        cuda_device = -1
    else:
        cuda_device = 0

    self.ee.cuda_device = cuda_device

    self.ee.elmo_bilm.to(device=flair.device)
    self.ee.elmo_bilm._elmo_lstm._states = tuple(
        [state.to(flair.device) for state in self.ee.elmo_bilm._elmo_lstm._states]
    )