Repository: flairNLP/flair

Issue:
classifier unable to make predictions when using elmo embeddings
Hi, I trained my text classifier using `ELMO pubmed` embeddings after training the model i am unable to predict on a gpu instance. The error and the code are mentioned below:-

Versions:-
```
allennlp== 0.9.0
flair== 0.5.1

```
code-

```
from flair.models import TextClassifier
from flair.data import Sentence
import pandas as pd 
classifier = TextClassifier.load('/content/drive/My Drive/best-model (2).pt')
data=pd.read_csv('/content/sample_prediction.csv')
data.head()
pred=[]

for index, sample in data.iterrows():
    sentence=Sentence(sample['sent'])
    classifier.predict(sentence)
    print(sentence.labels[0].value) 
    pred.append(sentence.labels[0].value)    

```

traceback:-
```

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
```
will really appreciate if someone can please help me out with this ? 


Discussion:
can someone please help me out here ?
I have the same problem
I hacked it by changing the line 88 in nn.py 

state = torch.load(f, map_location=flair.device) #'cpu')

from 'cpu' to 'flair.device' to force model load on GPU ...

it seems to be ok .. i am using this as a temporary measure until we get an official response :)

Resolving patch:
diff --git a/flair/embeddings/token.py b/flair/embeddings/token.py
--- a/flair/embeddings/token.py
+++ b/flair/embeddings/token.py
@@ -105,6 +105,7 @@ def get_named_embeddings_dict(self) -> Dict:
 
         return named_embeddings_dict
 
+
 class WordEmbeddings(TokenEmbeddings):
     """Standard static word embeddings, such as GloVe or FastText."""
 
@@ -1754,6 +1755,22 @@ def extra_repr(self):
     def __str__(self):
         return self.name
 
+    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])
+
 
 class NILCEmbeddings(WordEmbeddings):
     def __init__(self, embeddings: str, model: str = "skip", size: int = 100):


Test patch: