Merge pull request #52 from LiyuanLucasLiu/LiyuanLucasLiu-patch-2

upgrade to Pytorch0.4.1
LiyuanLucasLiu · Oct 3, 2018 · 0058bbd · 0058bbd
2 parents 533df89 + f575db5
commit 0058bbd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/model/predictor.py b/model/predictor.py
@@ -57,7 +57,7 @@ def decode_s(self, feature, label):
         current = None
 
         for f, y in zip(feature, label):
-            label = self.r_l_map[y]
+            label = self.r_l_map[y.item()]
 
             if label.startswith('B-'):
 
@@ -276,4 +276,4 @@ def apply_model(self, ner_model, features):
         scores = ner_model(f_f, f_p, b_f, b_p, w_f)
         decoded = self.decoder.decode(scores.data, mask_v)
 
-        return decoded
+        return decoded