Repository: flairNLP/flair

Issue:
Tests: Visualization test is failing on CUDA
Hi,

with the latest `master` version, one of the integration tests (`test_visual.py`) is failing when running on CUDA:

```bash
test_visual.py:90: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../flair/visual/activations.py:74: in highlight_selection
    rendered += self.highlight(activations[:, i], text)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <flair.visual.activations.Highlighter object at 0x7f207c682c88>
activation = tensor([ 1.2762e-02,  2.5522e-04,  7.5684e-03,  7.7306e-03, -1.5601e-02,
         6.5539e-07,  1.8706e-02,  6.5539e-07...4,  7.5684e-03,  7.7306e-03,  8.0798e-02,
        -1.1686e-03, -4.6697e-02,  7.7306e-03,  6.5539e-07], device='cuda:0')
text = 'The U.S. Centers for Disease Control and Prevention initially advised school systems to close if outbreaks occurred ,...f the virus meant most schools and day care centers should stay open , even if they had confirmed cases of swine flu .'

    def highlight(self, activation, text):
    
>       activation = activation.detach().numpy()
E       TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

../flair/visual/activations.py:37: TypeError
```


Discussion:
Thanks for reporting this - I'll take a look!

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -1522,10 +1522,10 @@ def embed(self, sentences: Union[List[Sentence], Sentence]):
             for add in range(longest_token_sequence_in_batch - len(sentence.tokens)):
                 word_embeddings.append(
                     torch.zeros(self.length_of_all_token_embeddings,
-                                dtype=torch.float, device=flair.device).unsqueeze(0)
+                                dtype=torch.float).unsqueeze(0)
                 )
 
-            word_embeddings_tensor = torch.cat(word_embeddings, 0)
+            word_embeddings_tensor = torch.cat(word_embeddings, 0).to(flair.device)
 
             sentence_states = word_embeddings_tensor
 
diff --git a/flair/models/text_classification_model.py b/flair/models/text_classification_model.py
--- a/flair/models/text_classification_model.py
+++ b/flair/models/text_classification_model.py
@@ -56,7 +56,7 @@ def forward(self, sentences) -> List[List[float]]:
         self.document_embeddings.embed(sentences)
 
         text_embedding_list = [sentence.get_embedding().unsqueeze(0) for sentence in sentences]
-        text_embedding_tensor = torch.cat(text_embedding_list, 0)
+        text_embedding_tensor = torch.cat(text_embedding_list, 0).to(flair.device)
 
         label_scores = self.decoder(text_embedding_tensor)
 
diff --git a/flair/visual/activations.py b/flair/visual/activations.py
--- a/flair/visual/activations.py
+++ b/flair/visual/activations.py
@@ -33,8 +33,7 @@ def __init__(self):
         ]
 
     def highlight(self, activation, text):
-
-        activation = activation.detach().numpy()
+        activation = activation.detach().cpu().numpy()
 
         step_size = (max(activation) - min(activation)) / len(self.color_map)
 


Test patch:
diff --git a/tests/test_data_fetchers.py b/tests/test_data_fetchers.py
--- a/tests/test_data_fetchers.py
+++ b/tests/test_data_fetchers.py
@@ -7,7 +7,7 @@
 
 def test_load_imdb_data(tasks_base_path):
     # get training, test and dev data
-    corpus = NLPTaskDataFetcher.load_corpus(NLPTask.IMDB, tasks_base_path)
+    corpus = NLPTaskDataFetcher.load_corpus('imdb', tasks_base_path)
 
     assert len(corpus.train) == 5
     assert len(corpus.dev) == 5
diff --git a/tests/test_hyperparameter.py b/tests/test_hyperparameter.py
--- a/tests/test_hyperparameter.py
+++ b/tests/test_hyperparameter.py
@@ -26,11 +26,11 @@ def test_sequence_tagger_param_selector(results_base_path, tasks_base_path):
     search_space.add(Parameter.DROPOUT, hp.uniform, low=0.25, high=0.75)
     search_space.add(Parameter.WORD_DROPOUT, hp.uniform, low=0.0, high=0.25)
     search_space.add(Parameter.LOCKED_DROPOUT, hp.uniform, low=0.0, high=0.5)
-    search_space.add(Parameter.HIDDEN_SIZE, hp.choice, options=[64, 128, 256, 512])
+    search_space.add(Parameter.HIDDEN_SIZE, hp.choice, options=[64, 128])
     search_space.add(Parameter.RNN_LAYERS, hp.choice, options=[1, 2])
 
     # model trainer parameter
-    search_space.add(Parameter.OPTIMIZER, hp.choice, options=[SGD, AdamW])
+    search_space.add(Parameter.OPTIMIZER, hp.choice, options=[SGD])
 
     # training parameter
     search_space.add(Parameter.MINI_BATCH_SIZE, hp.choice, options=[4, 8, 32])
@@ -49,7 +49,7 @@ def test_sequence_tagger_param_selector(results_base_path, tasks_base_path):
 
 @pytest.mark.integration
 def test_text_classifier_param_selector(results_base_path, tasks_base_path):
-    corpus = NLPTaskDataFetcher.load_corpus(NLPTask.IMDB, base_path=tasks_base_path)
+    corpus = NLPTaskDataFetcher.load_corpus('imdb', base_path=tasks_base_path)
 
     glove_embedding: WordEmbeddings = WordEmbeddings('en-glove')
 
@@ -60,7 +60,7 @@ def test_text_classifier_param_selector(results_base_path, tasks_base_path):
     search_space.add(Parameter.HIDDEN_SIZE, hp.choice, options=[64, 128, 256, 512])
     search_space.add(Parameter.RNN_LAYERS, hp.choice, options=[1, 2])
     search_space.add(Parameter.REPROJECT_WORDS, hp.choice, options=[True, False])
-    search_space.add(Parameter.REPROJECT_WORD_DIMENSION, hp.choice, options=[64, 128, 256, 512])
+    search_space.add(Parameter.REPROJECT_WORD_DIMENSION, hp.choice, options=[64, 128])
     search_space.add(Parameter.BIDIRECTIONAL, hp.choice, options=[True, False])
     search_space.add(Parameter.DROPOUT, hp.uniform, low=0.25, high=0.75)
     search_space.add(Parameter.WORD_DROPOUT, hp.uniform, low=0.25, high=0.75)
diff --git a/tests/test_model_integration.py b/tests/test_model_integration.py
--- a/tests/test_model_integration.py
+++ b/tests/test_model_integration.py
@@ -1,6 +1,7 @@
 import os
 import shutil
 import pytest
+from torch.optim import SGD
 
 from torch.optim.optimizer import Optimizer
 from torch.optim.adam import Adam
@@ -23,7 +24,7 @@ def test_train_load_use_tagger(results_base_path, tasks_base_path):
 
     embeddings = WordEmbeddings('glove')
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',
@@ -56,7 +57,7 @@ def test_train_load_use_tagger_large(results_base_path, tasks_base_path):
 
     embeddings = WordEmbeddings('glove')
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='pos',
@@ -89,7 +90,7 @@ def test_train_charlm_load_use_tagger(results_base_path, tasks_base_path):
 
     embeddings = FlairEmbeddings('news-forward-fast')
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',
@@ -125,7 +126,7 @@ def test_train_charlm_changed_chache_load_use_tagger(results_base_path, tasks_ba
     os.makedirs(cache_dir, exist_ok=True)
     embeddings = FlairEmbeddings('news-forward-fast', cache_directory=cache_dir)
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',
@@ -161,7 +162,7 @@ def test_train_charlm_nochache_load_use_tagger(results_base_path, tasks_base_pat
 
     embeddings = FlairEmbeddings('news-forward-fast', use_cache=False)
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',
@@ -194,7 +195,7 @@ def test_train_optimizer(results_base_path, tasks_base_path):
 
     embeddings = WordEmbeddings('glove')
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',
@@ -230,7 +231,7 @@ def test_train_optimizer_arguments(results_base_path, tasks_base_path):
 
     embeddings = WordEmbeddings('glove')
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',
@@ -266,13 +267,13 @@ def test_find_learning_rate(results_base_path, tasks_base_path):
 
     embeddings = WordEmbeddings('glove')
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',
                                             use_crf=False)
 
-    optimizer: Optimizer = AdamW
+    optimizer: Optimizer = SGD
 
     # initialize trainer
     trainer: ModelTrainer = ModelTrainer(tagger, corpus, optimizer=optimizer)
@@ -308,7 +309,7 @@ def test_load_use_serialized_tagger():
 @pytest.mark.integration
 def test_train_load_use_classifier(results_base_path, tasks_base_path):
 
-    corpus = NLPTaskDataFetcher.load_corpus(NLPTask.IMDB, base_path=tasks_base_path)
+    corpus = NLPTaskDataFetcher.load_corpus('imdb', base_path=tasks_base_path)
     label_dict = corpus.make_label_dictionary()
 
     glove_embedding: WordEmbeddings = WordEmbeddings('en-glove')
@@ -344,7 +345,7 @@ def test_train_load_use_classifier(results_base_path, tasks_base_path):
 @pytest.mark.integration
 def test_train_charlm_load_use_classifier(results_base_path, tasks_base_path):
 
-    corpus = NLPTaskDataFetcher.load_corpus(NLPTask.IMDB, base_path=tasks_base_path)
+    corpus = NLPTaskDataFetcher.load_corpus('imdb', base_path=tasks_base_path)
     label_dict = corpus.make_label_dictionary()
 
     glove_embedding: TokenEmbeddings = FlairEmbeddings('news-forward-fast')
@@ -380,7 +381,7 @@ def test_train_charlm_load_use_classifier(results_base_path, tasks_base_path):
 @pytest.mark.integration
 def test_train_charlm_nocache_load_use_classifier(results_base_path, tasks_base_path):
 
-    corpus = NLPTaskDataFetcher.load_corpus(NLPTask.IMDB, base_path=tasks_base_path)
+    corpus = NLPTaskDataFetcher.load_corpus('imdb', base_path=tasks_base_path)
     label_dict = corpus.make_label_dictionary()
 
     glove_embedding: TokenEmbeddings = FlairEmbeddings('news-forward-fast', use_cache=False)
@@ -453,7 +454,7 @@ def test_train_load_use_tagger_multicorpus(results_base_path, tasks_base_path):
 
     embeddings = WordEmbeddings('glove')
 
-    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+    tagger: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',
@@ -479,7 +480,7 @@ def test_train_load_use_tagger_multicorpus(results_base_path, tasks_base_path):
 
 @pytest.mark.integration
 def test_train_resume_text_classification_training(results_base_path, tasks_base_path):
-    corpus = NLPTaskDataFetcher.load_corpus(NLPTask.IMDB, base_path=tasks_base_path)
+    corpus = NLPTaskDataFetcher.load_corpus('imdb', base_path=tasks_base_path)
     label_dict = corpus.make_label_dictionary()
 
     embeddings: TokenEmbeddings = FlairEmbeddings('news-forward-fast', use_cache=False)
@@ -504,7 +505,7 @@ def test_train_resume_sequence_tagging_training(results_base_path, tasks_base_pa
 
     embeddings = WordEmbeddings('glove')
 
-    model: SequenceTagger = SequenceTagger(hidden_size=256,
+    model: SequenceTagger = SequenceTagger(hidden_size=64,
                                             embeddings=embeddings,
                                             tag_dictionary=tag_dictionary,
                                             tag_type='ner',