commit

run_segbot.py

@@ -11,13 +11,6 @@ import MeCab
 import pysbd
 
 
-class CPU_Unpickler(pickle.Unpickler):
-    def find_class(self, module, name):
-        if module == 'torch.storage' and name == '_load_from_bytes':
-            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
-        else: return super().find_class(module, name)
-
-
 def create_data(doc,fm,split_method):
     wakati = MeCab.Tagger("-Owakati -b 81920")
     seg = pysbd.Segmenter(language="ja", clean=False)
@@ -74,7 +67,7 @@ def setup():
     with open('index2word.pickle', 'rb') as f:    
         index2word = pickle.load(f)
     with open('model.pickle', 'rb') as f:
-        mysolver = CPU_Unpickler(f).load()
+        mysolver = torch.load(io.BytesIO(f), map_location='cpu')
     with open('fm.pickle', 'rb') as f:
         fm = pickle.load(f)