Upload model_manager.py

common/model_manager.py

@@ -2,7 +2,7 @@
 Author: Qiguang Chen
 Date: 2023-01-11 10:39:26
 LastEditors: Qiguang Chen
-LastEditTime: 2023-02-08 01:02:20
+LastEditTime: 2023-02-08 00:57:09
 Description: manage all process of model training and prediction.
 
 '''
@@ -287,8 +287,7 @@ class ModelManager(object):
         return outputs, res
 
     def load(self):
-        # self.model = torch.load(os.path.join(self.config.base["model_dir"], "model.pkl"), map_location=torch.device(self.device))
-        self.model = torch.load(os.path.join(self.config.base["model_dir"], "model.pkl"), map_location=torch.device("cpu"))
+        self.model = torch.load(os.path.join(self.config.base["model_dir"], "model.pkl"), map_location=torch.device(self.device))
         if self.config.tokenizer["_tokenizer_name_"] == "word_tokenizer":
             self.tokenizer = get_tokenizer_class(self.config.tokenizer["_tokenizer_name_"]).from_file(
                 os.path.join(self.config.base["model_dir"], "tokenizer.json"))