Repository: flairNLP/flair

Issue:
aggregated_embedding not working with CPU
**Describe the bug**
error in the end of the learning for 90 epoch using the other language word embeddings

**To Reproduce**
```python
    WordEmbeddings('ar'),

    # contextual string embeddings, forward
    PooledFlairEmbeddings('ar-forward', pooling='min'),

    # contextual string embeddings, backward
    PooledFlairEmbeddings('ar-backward', pooling='min'),
```

**Expected behavior**
A clear and concise description of what you expected to happen.

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environment (please complete the following information):**
 - OS [e.g. iOS, Linux]:
 - Version [e.g. flair-0.3.2]:

**Additional context**
```console
2020-02-03 16:59:38,777 Reading data from /content/gdrive/My Drive/resources/tasks/conll_03
2020-02-03 16:59:38,780 Train: /content/gdrive/My Drive/resources/tasks/conll_03/train.txt
2020-02-03 16:59:38,782 Dev: /content/gdrive/My Drive/resources/tasks/conll_03/dev.txt
2020-02-03 16:59:38,785 Test: /content/gdrive/My Drive/resources/tasks/conll_03/test.txt
2020-02-03 17:00:24,966 ----------------------------------------------------------------------------------------------------
2020-02-03 17:00:24,970 Model: "SequenceTagger(
  (embeddings): StackedEmbeddings(
    (list_embedding_0): WordEmbeddings('ar')
    (list_embedding_1): PooledFlairEmbeddings(
      (context_embeddings): FlairEmbeddings(
        (lm): LanguageModel(
          (drop): Dropout(p=0.1, inplace=False)
          (encoder): Embedding(7125, 100)
          (rnn): LSTM(100, 2048)
          (decoder): Linear(in_features=2048, out_features=7125, bias=True)
        )
      )
    )
    (list_embedding_2): PooledFlairEmbeddings(
      (context_embeddings): FlairEmbeddings(
        (lm): LanguageModel(
          (drop): Dropout(p=0.1, inplace=False)
          (encoder): Embedding(7125, 100)
          (rnn): LSTM(100, 2048)
          (decoder): Linear(in_features=2048, out_features=7125, bias=True)
        )
      )
    )
  )
  (word_dropout): WordDropout(p=0.05)
  (locked_dropout): LockedDropout(p=0.5)
  (embedding2nn): Linear(in_features=8492, out_features=8492, bias=True)
  (rnn): LSTM(8492, 256, batch_first=True, bidirectional=True)
  (linear): Linear(in_features=512, out_features=12, bias=True)
  (beta): 1.0
  (weights): None
  (weight_tensor) None
)"
2020-02-03 17:00:24,975 ----------------------------------------------------------------------------------------------------
2020-02-03 17:00:24,978 Corpus: "Corpus: 1328 train + 710 dev + 605 test sentences"
2020-02-03 17:00:24,981 ----------------------------------------------------------------------------------------------------
2020-02-03 17:00:24,985 Parameters:
2020-02-03 17:00:24,987  - learning_rate: "0.1"
2020-02-03 17:00:24,990  - mini_batch_size: "32"
2020-02-03 17:00:24,991  - patience: "3"
2020-02-03 17:00:24,994  - anneal_factor: "0.5"
2020-02-03 17:00:24,996  - max_epochs: "90"
2020-02-03 17:00:24,999  - shuffle: "True"
2020-02-03 17:00:25,001  - train_with_dev: "True"
2020-02-03 17:00:25,002  - batch_growth_annealing: "False"
2020-02-03 17:00:25,004 ----------------------------------------------------------------------------------------------------
2020-02-03 17:00:25,007 Model training base path: "/content/gdrive/My Drive/resources/taggers/example-ner"
2020-02-03 17:00:25,011 ----------------------------------------------------------------------------------------------------
2020-02-03 17:00:25,014 Device: cuda:0
2020-02-03 17:00:25,015 ----------------------------------------------------------------------------------------------------
2020-02-03 17:00:25,020 Embeddings storage mode: cpu
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type StackedEmbeddings. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type WordEmbeddings. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type PooledFlairEmbeddings. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type FlairEmbeddings. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type LanguageModel. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type Dropout. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type Embedding. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type LSTM. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
/usr/local/lib/python3.6/dist-packages/torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type Linear. It won't be checked for correctness upon loading.
  "type " + obj.__name__ + ". It won't be checked "
2020-02-03 17:00:45,361 ----------------------------------------------------------------------------------------------------
2020-02-03 17:00:45,365 Testing using best model ...
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-10-4c66a380caa5> in <module>()
     35               mini_batch_size=32,
     36               max_epochs=90,
---> 37               checkpoint=True)

6 frames
/usr/local/lib/python3.6/dist-packages/flair/embeddings.py in _add_embeddings_internal(self, sentences)
   2155                         else:
   2156                             aggregated_embedding = self.aggregate_op(
-> 2157                                 self.word_embeddings[token.text], local_embedding
   2158                             )
   2159                             if self.pooling == "fade":

RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'other' in call to _th_min
```


Discussion:

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -2180,6 +2180,13 @@ def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
     def embedding_length(self) -> int:
         return self.embedding_length
 
+    def __setstate__(self, d):
+        self.__dict__ = d
+
+        if flair.device != 'cpu':
+            for key in self.word_embeddings:
+                self.word_embeddings[key] = self.word_embeddings[key].cpu()
+
 
 class BertEmbeddings(TokenEmbeddings):
     def __init__(
@@ -2999,7 +3006,7 @@ def _add_embeddings_internal(self, sentences: Union[List[Sentence], Sentence]):
 
     def _apply(self, fn):
         major, minor, build, *_ = (int(info)
-                                for info in torch.__version__.replace("+cpu","").split('.'))
+                                for info in torch.__version__.replace("+",".").split('.') if info.isdigit())
 
         # fixed RNN change format for torch 1.4.0
         if major >= 1 and minor >= 4:


Test patch: