fixed wikisql data loading

pcyin · Mar 2, 2019 · 5c715f4 · 5c715f4
1 parent 0f8fbca
commit 5c715f4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/model/parser.py b/model/parser.py
@@ -796,8 +796,8 @@ def save(self, path):
         }
         torch.save(params, path)
 
-    @staticmethod
-    def load(model_path, cuda=False):
+    @classmethod
+    def load(cls, model_path, cuda=False):
         params = torch.load(model_path, map_location=lambda storage, loc: storage)
         vocab = params['vocab']
         transition_system = params['transition_system']
@@ -807,7 +807,7 @@ def load(model_path, cuda=False):
         saved_state = params['state_dict']
         saved_args.cuda = cuda
 
-        parser = Parser(saved_args, vocab, transition_system)
+        parser = cls(saved_args, vocab, transition_system)
 
         parser.load_state_dict(saved_state)