Merge pull request #42 from lsvih/patch-1

Fix fastHan model loader in cpu-only env

Author: fdugzc

fastHan/FastModel.py

@@ -161,7 +161,10 @@ def __init__(self, model_type='base', url=None):
             task_label_map,
             self.all_tasks,
             ensembledWeightManager=self.ensembledWeightManager)
-        self.model.load_state_dict(torch.load(model_path))
+        if torch.cuda.is_available():
+            self.model.load_state_dict(torch.load(model_path))
+        else:
+            self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
         self.model.to(self.device)
         self.model.eval()
         # 加载tokenizer
@@ -653,4 +656,4 @@ def __call__(self,
         answer = []
         for ans_ in ans:
             answer.append(Sentence(ans_, target))
-        return answer
+        return answer