fix cpu

model.py

@@ -173,6 +173,11 @@ class Model(torch.jit.ScriptModule):
         return path_to_checkpoint_file
 
     def restore(self, path_to_checkpoint_file):
-        self.load_state_dict(torch.load(path_to_checkpoint_file))
+        self.load_state_dict(
+            torch.load(
+                path_to_checkpoint_file,
+                map_location=torch.device('cpu')
+            )
+        )
         step = int(path_to_checkpoint_file.split('\\')[-1][6:-4])
         return step