GH-407: fix training on cuda without CRF

flairNLP · Feb 2, 2019 · abc05ef · abc05ef
1 parent 1876f1f
commit abc05ef
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/flair/models/sequence_tagger_model.py b/flair/models/sequence_tagger_model.py
@@ -429,10 +429,11 @@ def _calculate_loss(self, features, lengths, tags) -> float:
             for sentence_feats, sentence_tags, sentence_length in zip(features, tags, lengths):
                 sentence_feats = sentence_feats[:sentence_length]
 
-                tag_tensor = torch.LongTensor(sentence_tags)
-                tag_tensor = tag_tensor.to(flair.device)
+                # print(sentence_tags)
+                # tag_tensor = torch.LongTensor(sentence_tags)
+                # tag_tensor = tag_tensor.to(flair.device)
 
-                score += torch.nn.functional.cross_entropy(sentence_feats, tag_tensor)
+                score += torch.nn.functional.cross_entropy(sentence_feats, sentence_tags)
 
             return score