Update (#3916)

dglgo/README.md

@@ -319,7 +319,7 @@ def main():
         **pipeline_cfg["optimizer"])
     # train
     test_acc = train(cfg, pipeline_cfg, device, data, model, optimizer, loss)
-    torch.save(model, pipeline_cfg["save_path"])
+    torch.save(model.state_dict(), pipeline_cfg["save_path"])
     return test_acc
 
 ...

dglgo/dglgo/pipeline/linkpred/linkpred.jinja-py

@@ -112,7 +112,7 @@ def main():
     loss = torch.nn.{{ loss }}()
     optimizer = torch.optim.Adam(params, **pipeline_cfg["optimizer"])
     test_hits = train(cfg, pipeline_cfg, device, dataset, model, optimizer, loss)
-    torch.save(model, pipeline_cfg["save_path"])
+    torch.save(model.state_dict(), pipeline_cfg["save_path"])
     return test_hits
 
 if __name__ == '__main__':

dglgo/dglgo/pipeline/nodepred/nodepred.jinja-py

@@ -112,7 +112,7 @@ def main():
     optimizer = torch.optim.{{ user_cfg.general_pipeline.optimizer.name }}(model.parameters(), **pipeline_cfg["optimizer"])
     # train
     test_acc = train(cfg, pipeline_cfg, device, data, model, optimizer, loss)
-    torch.save(model, pipeline_cfg["save_path"])
+    torch.save(model.state_dict(), pipeline_cfg["save_path"])
     return test_acc
 
 if __name__ == '__main__':

dglgo/dglgo/pipeline/nodepred_sample/nodepred-ns.jinja-py

@@ -158,7 +158,7 @@ def main():
     loss = torch.nn.{{ user_cfg.general_pipeline.loss }}()
     optimizer = torch.optim.{{ user_cfg.general_pipeline.optimizer.name }}(model.parameters(), **pipeline_cfg["optimizer"])
     test_acc = train(cfg, pipeline_cfg, device, data, model, optimizer, loss)
-    torch.save(model, pipeline_cfg["save_path"])
+    torch.save(model.state_dict(), pipeline_cfg["save_path"])
     return test_acc
 
 if __name__ == '__main__':