commit

run_segbot.py

@@ -10,6 +10,14 @@ import gensim
 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)
@@ -66,7 +74,7 @@ def setup():
     with open('index2word.pickle', 'rb') as f:    
         index2word = pickle.load(f)
     with open('model.pickle', 'rb') as f:
-        mysolver = torch.load(f, torch.device('cpu'))
+        mysolver = CPU_Unpickler(f).load()
     with open('fm.pickle', 'rb') as f:
         fm = pickle.load(f)