change pkl to pt

HazyResearch · May 4, 2020 · f846be3 · f846be3
1 parent d9b3d00
commit f846be3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/run.py b/run.py
@@ -157,7 +157,7 @@ def train(args):
                 counter = 0
                 best_epoch = step
                 logging.info("\t Saving model at epoch {} in {}".format(step, save_dir))
-                torch.save(model.cpu().state_dict(), os.path.join(save_dir, "model.pkl"))
+                torch.save(model.cpu().state_dict(), os.path.join(save_dir, "model.pt"))
                 model.cuda()
             else:
                 counter += 1
@@ -171,10 +171,10 @@ def train(args):
 
     logging.info("\t Optimization finished")
     if not best_mrr:
-        torch.save(model.cpu().state_dict(), os.path.join(save_dir, "model.pkl"))
+        torch.save(model.cpu().state_dict(), os.path.join(save_dir, "model.pt"))
     else:
         logging.info("\t Loading best model saved at epoch {}".format(best_epoch))
-        model.load_state_dict(torch.load(os.path.join(save_dir, "model.pkl")))
+        model.load_state_dict(torch.load(os.path.join(save_dir, "model.pt")))
     model.cuda()
     model.eval()
 

diff --git a/test.py b/test.py
@@ -32,7 +32,7 @@ def test(model_dir):
     model = getattr(models, args.model)(args)
     device = 'cuda'
     model.to(device)
-    model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pkl')))
+    model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pt')))
 
     # eval
     test_metrics = avg_both(*model.compute_metrics(test_examples, filters))