[TestFailure] Resolving exceptions thrown across multiple GraphBolt tests. (dmlc#7852)

drivanov · web-flow · commit 540dd2ba4d21 · 2025-01-08T14:05:34.000+08:00
diff --git a/dglgo/dglgo/apply_pipeline/graphpred/gen.py b/dglgo/dglgo/apply_pipeline/graphpred/gen.py
@@ -63,7 +63,7 @@ def config(
             cpt: str = typer.Option(..., help="input checkpoint file path"),
         ):
             # Training configuration
-            train_cfg = torch.load(cpt)["cfg"]
+            train_cfg = torch.load(cpt, weights_only=False)["cfg"]
             if data is None:
                 print("data is not specified, use the training dataset")
                 data = train_cfg["data_name"]
@@ -119,7 +119,9 @@ def gen_script(cls, user_cfg_dict):
         cls.user_cfg_cls(**user_cfg_dict)
 
         # Training configuration
-        train_cfg = torch.load(user_cfg_dict["cpt_path"])["cfg"]
+        train_cfg = torch.load(user_cfg_dict["cpt_path"], weights_only=False)[
+            "cfg"
+        ]
 
         # Dict for code rendering
         render_cfg = deepcopy(user_cfg_dict)
diff --git a/dglgo/dglgo/apply_pipeline/graphpred/graphpred.jinja-py b/dglgo/dglgo/apply_pipeline/graphpred/graphpred.jinja-py
@@ -57,7 +57,7 @@ def main():
         data to have the same number of input edge features, got {:d} and {:d}'.format(model_edge_feat_size, data_edge_feat_size)
 
     model = {{ model_class_name }}(**cfg['model'])
-    model.load_state_dict(torch.load(cfg['cpt_path'], map_location='cpu')['model'])
+    model.load_state_dict(torch.load(cfg['cpt_path'], weights_only=False, map_location='cpu')['model'])
     pred = infer(device, data_loader, model).detach().cpu()
 
     # Dump the results
diff --git a/dglgo/dglgo/apply_pipeline/nodepred/gen.py b/dglgo/dglgo/apply_pipeline/nodepred/gen.py
@@ -48,7 +48,7 @@ def config(
             cpt: str = typer.Option(..., help="input checkpoint file path"),
         ):
             # Training configuration
-            train_cfg = torch.load(cpt)["cfg"]
+            train_cfg = torch.load(cpt, weights_only=False)["cfg"]
             if data is None:
                 print("data is not specified, use the training dataset")
                 data = train_cfg["data_name"]
@@ -101,7 +101,9 @@ def gen_script(cls, user_cfg_dict):
         cls.user_cfg_cls(**user_cfg_dict)
 
         # Training configuration
-        train_cfg = torch.load(user_cfg_dict["cpt_path"])["cfg"]
+        train_cfg = torch.load(user_cfg_dict["cpt_path"], weights_only=False)[
+            "cfg"
+        ]
 
         # Dict for code rendering
         render_cfg = deepcopy(user_cfg_dict)
diff --git a/dglgo/dglgo/apply_pipeline/nodepred/nodepred.jinja-py b/dglgo/dglgo/apply_pipeline/nodepred/nodepred.jinja-py
@@ -48,7 +48,7 @@ def main():
                 features, got {:d} and {:d}'.format(model_in_size, data_in_size)
 
     model = {{ model_class_name }}(**cfg['model'])
-    model.load_state_dict(torch.load(cfg['cpt_path'], map_location='cpu')['model'])
+    model.load_state_dict(torch.load(cfg['cpt_path'], weights_only=False, map_location='cpu')['model'])
     logits = infer(device, data, model)
     pred = logits.argmax(dim=1).cpu()
 
diff --git a/dglgo/dglgo/apply_pipeline/nodepred_sample/gen.py b/dglgo/dglgo/apply_pipeline/nodepred_sample/gen.py
@@ -48,7 +48,7 @@ def config(
             cpt: str = typer.Option(..., help="input checkpoint file path"),
         ):
             # Training configuration
-            train_cfg = torch.load(cpt)["cfg"]
+            train_cfg = torch.load(cpt, weights_only=False)["cfg"]
             if data is None:
                 print("data is not specified, use the training dataset")
                 data = train_cfg["data_name"]
@@ -101,7 +101,9 @@ def gen_script(cls, user_cfg_dict):
         cls.user_cfg_cls(**user_cfg_dict)
 
         # Training configuration
-        train_cfg = torch.load(user_cfg_dict["cpt_path"])["cfg"]
+        train_cfg = torch.load(user_cfg_dict["cpt_path"], weights_only=False)[
+            "cfg"
+        ]
 
         # Dict for code rendering
         render_cfg = deepcopy(user_cfg_dict)
diff --git a/dglgo/dglgo/apply_pipeline/nodepred_sample/nodepred-ns.jinja-py b/dglgo/dglgo/apply_pipeline/nodepred_sample/nodepred-ns.jinja-py
@@ -48,7 +48,7 @@ def main():
                 features, got {:d} and {:d}'.format(model_in_size, data_in_size)
 
     model = {{ model_class_name }}(**cfg['model'])
-    model.load_state_dict(torch.load(cfg['cpt_path'], map_location='cpu')['model'])
+    model.load_state_dict(torch.load(cfg['cpt_path'], weights_only=False, map_location='cpu')['model'])
     logits = infer(device, data, model)
     pred = logits.argmax(dim=1).cpu()
 
diff --git a/dglgo/dglgo/pipeline/graphpred/graphpred.jinja-py b/dglgo/dglgo/pipeline/graphpred/graphpred.jinja-py
@@ -101,7 +101,7 @@ def main(run, cfg, data):
         else:
             lr_scheduler.step()
 
-    model.load_state_dict(torch.load(tmp_cpt_path))
+    model.load_state_dict(torch.load(tmp_cpt_path, weights_only=False))
     os.remove(tmp_cpt_path)
     test_metric = evaluate(device, test_loader, model)
     print('Test Metric: {:.4f}'.format(test_metric))
diff --git a/dglgo/dglgo/pipeline/nodepred/nodepred.jinja-py b/dglgo/dglgo/pipeline/nodepred/nodepred.jinja-py
@@ -42,7 +42,7 @@ class EarlyStopping:
         torch.save(model.state_dict(), self.checkpoint_path)
 
     def load_checkpoint(self, model):
-        model.load_state_dict(torch.load(self.checkpoint_path))
+        model.load_state_dict(torch.load(self.checkpoint_path, weights_only=False))
 
     def close(self):
         os.remove(self.checkpoint_path)
diff --git a/dglgo/dglgo/pipeline/nodepred_sample/nodepred-ns.jinja-py b/dglgo/dglgo/pipeline/nodepred_sample/nodepred-ns.jinja-py
@@ -42,7 +42,7 @@ class EarlyStopping:
         torch.save(model.state_dict(), self.checkpoint_path)
 
     def load_checkpoint(self, model):
-        model.load_state_dict(torch.load(self.checkpoint_path))
+        model.load_state_dict(torch.load(self.checkpoint_path, weights_only=False))
 
     def close(self):
         os.remove(self.checkpoint_path)
diff --git a/dglgo/dglgo/utils/early_stop.py b/dglgo/dglgo/utils/early_stop.py
@@ -34,4 +34,6 @@ def save_checkpoint(self, model):
         torch.save(model.state_dict(), self.checkpoint_path)
 
     def load_checkpoint(self, model):
-        model.load_state_dict(torch.load(self.checkpoint_path))
+        model.load_state_dict(
+            torch.load(self.checkpoint_path, weights_only=False)
+        )
diff --git a/examples/pytorch/GNN-FiLM/main.py b/examples/pytorch/GNN-FiLM/main.py
@@ -194,7 +194,9 @@ def main(args):
     model.eval()
     test_loss = []
     test_f1 = []
-    model.load_state_dict(torch.load(os.path.join(args.save_dir, args.name)))
+    model.load_state_dict(
+        torch.load(os.path.join(args.save_dir, args.name), weights_only=False)
+    )
     with torch.no_grad():
         for batch in test_set:
             g = batch.graph
diff --git a/examples/pytorch/TAHIN/main.py b/examples/pytorch/TAHIN/main.py
@@ -111,7 +111,9 @@ def main(args):
     # test use the best model
     model.eval()
     with torch.no_grad():
-        model.load_state_dict(torch.load("TAHIN" + "_" + args.dataset))
+        model.load_state_dict(
+            torch.load("TAHIN" + "_" + args.dataset, weights_only=False)
+        )
         test_loss = []
         test_acc = []
         test_auc = []
diff --git a/examples/pytorch/argo/main.py b/examples/pytorch/argo/main.py
@@ -180,7 +180,7 @@ def train(
 
     PATH = "model.pt"
     if counter[0] != 0:
-        checkpoint = torch.load(PATH)
+        checkpoint = torch.load(PATH, weights_only=False)
         model.load_state_dict(checkpoint["model_state_dict"])
         opt.load_state_dict(checkpoint["optimizer_state_dict"])
         epoch = checkpoint["epoch"]
diff --git a/examples/pytorch/correct_and_smooth/main.py b/examples/pytorch/correct_and_smooth/main.py
@@ -66,7 +66,9 @@ def main():
     if args.pretrain:
         print("---------- Before ----------")
         model.load_state_dict(
-            torch.load(f"base/{args.dataset}-{args.model}.pt")
+            torch.load(
+                f"base/{args.dataset}-{args.model}.pt", weights_only=False
+            )
         )
         model.eval()
 
diff --git a/examples/pytorch/dgi/train.py b/examples/pytorch/dgi/train.py
@@ -126,7 +126,7 @@ def main(args):
 
     # train classifier
     print("Loading {}th epoch".format(best_t))
-    dgi.load_state_dict(torch.load("best_dgi.pkl"))
+    dgi.load_state_dict(torch.load("best_dgi.pkl", weights_only=False))
     embeds = dgi.encoder(features, corrupt=False)
     embeds = embeds.detach()
     mean = 0
diff --git a/examples/pytorch/diffpool/train.py b/examples/pytorch/diffpool/train.py
@@ -223,7 +223,8 @@ def graph_classify_task(prog_args):
                 + "/"
                 + prog_args.dataset
                 + "/model.iter-"
-                + str(prog_args.load_epoch)
+                + str(prog_args.load_epoch),
+                weights_only=False,
             )
         )
 
@@ -334,7 +335,8 @@ def evaluate(dataloader, model, prog_args, logger=None):
                 + "/"
                 + prog_args.dataset
                 + "/model.iter-"
-                + str(logger["best_epoch"])
+                + str(logger["best_epoch"]),
+                weights_only=False,
             )
         )
     model.eval()
diff --git a/examples/pytorch/dimenet/main.py b/examples/pytorch/dimenet/main.py
@@ -238,7 +238,9 @@ def main(model_cnf):
     if pretrain_params["flag"]:
         torch_path = pretrain_params["path"]
         target = model_params["targets"][0]
-        model.load_state_dict(torch.load(f"{torch_path}/{target}.pt"))
+        model.load_state_dict(
+            torch.load(f"{torch_path}/{target}.pt", weights_only=False)
+        )
 
         logger.info("Testing with Pretrained model")
         predictions, labels = evaluate(device, model, test_loader)
diff --git a/examples/pytorch/gatv2/train.py b/examples/pytorch/gatv2/train.py
@@ -178,7 +178,9 @@ def main(args):
 
     print()
     if args.early_stop:
-        model.load_state_dict(torch.load("es_checkpoint.pt"))
+        model.load_state_dict(
+            torch.load("es_checkpoint.pt", weights_only=False)
+        )
     acc = evaluate(g, model, features, labels, test_mask)
     print("Test Accuracy {:.4f}".format(acc))
 
diff --git a/examples/pytorch/graphsaint/train_sampling.py b/examples/pytorch/graphsaint/train_sampling.py
@@ -214,7 +214,10 @@ def main(args, task):
     # test
     if args.use_val:
         model.load_state_dict(
-            torch.load(os.path.join(log_dir, "best_model_{}.pkl".format(task)))
+            torch.load(
+                os.path.join(log_dir, "best_model_{}.pkl".format(task)),
+                weights_only=False,
+            )
         )
     if cpu_flag and cuda:
         model = model.to("cpu")
diff --git a/examples/pytorch/graphwriter/train.py b/examples/pytorch/graphwriter/train.py
@@ -161,7 +161,7 @@ def main(args):
     model = GraphWriter(args)
     model.to(args.device)
     if args.test:
-        model = torch.load(args.save_model)
+        model = torch.load(args.save_model, weights_only=False)
         model.args = args
         print(model)
         test(model, test_dataloader, args)
diff --git a/examples/pytorch/han/utils.py b/examples/pytorch/han/utils.py
@@ -304,4 +304,4 @@ def save_checkpoint(self, model):
 
     def load_checkpoint(self, model):
         """Load the latest checkpoint."""
-        model.load_state_dict(torch.load(self.filename))
+        model.load_state_dict(torch.load(self.filename, weights_only=False))
diff --git a/examples/pytorch/hardgat/train.py b/examples/pytorch/hardgat/train.py
@@ -154,7 +154,9 @@ def main(args):
 
     print()
     if args.early_stop:
-        model.load_state_dict(torch.load("es_checkpoint.pt"))
+        model.load_state_dict(
+            torch.load("es_checkpoint.pt", weights_only=False)
+        )
     acc = evaluate(model, features, labels, test_mask)
     print("Test Accuracy {:.4f}".format(acc))
 
diff --git a/examples/pytorch/hilander/PSS/Smooth_AP/src/netlib.py b/examples/pytorch/hilander/PSS/Smooth_AP/src/netlib.py
@@ -71,7 +71,9 @@ def networkselect(opt):
         raise Exception("Network {} not available!".format(opt.arch))
 
     if opt.resume:
-        weights = torch.load(os.path.join(opt.save_path, opt.resume))
+        weights = torch.load(
+            os.path.join(opt.save_path, opt.resume), weights_only=False
+        )
         weights_state_dict = weights["state_dict"]
 
         if torch.cuda.device_count() > 1:
diff --git a/examples/pytorch/hilander/PSS/test_subg_inat.py b/examples/pytorch/hilander/PSS/test_subg_inat.py
@@ -173,7 +173,7 @@
         use_cluster_feat=args.use_cluster_feat,
         use_focal_loss=args.use_focal_loss,
     )
-    model.load_state_dict(torch.load(args.model_filename))
+    model.load_state_dict(torch.load(args.model_filename, weights_only=False))
     model = model.to(device)
     model.eval()
 
diff --git a/examples/pytorch/hilander/test.py b/examples/pytorch/hilander/test.py
@@ -83,7 +83,7 @@
         use_cluster_feat=args.use_cluster_feat,
         use_focal_loss=args.use_focal_loss,
     )
-    model.load_state_dict(torch.load(args.model_filename))
+    model.load_state_dict(torch.load(args.model_filename, weights_only=False))
     model = model.to(device)
     model.eval()
 
diff --git a/examples/pytorch/hilander/test_subg.py b/examples/pytorch/hilander/test_subg.py
@@ -104,7 +104,7 @@
         use_cluster_feat=args.use_cluster_feat,
         use_focal_loss=args.use_focal_loss,
     )
-    model.load_state_dict(torch.load(args.model_filename))
+    model.load_state_dict(torch.load(args.model_filename, weights_only=False))
     model = model.to(device)
     model.eval()
 
diff --git a/examples/pytorch/jtnn/vaetrain_dgl.py b/examples/pytorch/jtnn/vaetrain_dgl.py
@@ -54,7 +54,7 @@ def worker_init_fn(id_):
 model = DGLJTNNVAE(vocab, hidden_size, latent_size, depth)
 
 if opts.model_path is not None:
-    model.load_state_dict(torch.load(opts.model_path))
+    model.load_state_dict(torch.load(opts.model_path, weights_only=False))
 else:
     for param in model.parameters():
         if param.dim() == 1:
diff --git a/examples/pytorch/lda/lda_model.py b/examples/pytorch/lda/lda_model.py
@@ -496,6 +496,6 @@ def doc_subgraph(G, doc_ids):
     with io.BytesIO() as f:
         model.save(f)
         f.seek(0)
-        print(torch.load(f))
+        print(torch.load(f, weights_only=False))
 
     print("Testing LatentDirichletAllocation passed!")
diff --git a/examples/pytorch/ogb/ngnn_seal/main.py b/examples/pytorch/ogb/ngnn_seal/main.py
@@ -625,8 +625,12 @@ def print_log(*x, sep="\n", end="\n", mode="a"):
                         args.res_dir,
                         f"run{run+1}_optimizer_checkpoint{epoch}.pth",
                     )
-                    model.load_state_dict(torch.load(model_name))
-                    optimizer.load_state_dict(torch.load(optimizer_name))
+                    model.load_state_dict(
+                        torch.load(model_name, weights_only=False)
+                    )
+                    optimizer.load_state_dict(
+                        torch.load(optimizer_name, weights_only=False)
+                    )
                     tested[epoch] = (
                         test(final_val_loader, dataset.eval_metric)[
                             dataset.eval_metric
diff --git a/examples/pytorch/ogb/ogbn-arxiv/correct_and_smooth.py b/examples/pytorch/ogb/ogbn-arxiv/correct_and_smooth.py
@@ -179,7 +179,7 @@ def main():
 
     for pred_file in glob.iglob(args.pred_files):
         print("load:", pred_file)
-        pred = torch.load(pred_file)
+        pred = torch.load(pred_file, weights_only=False)
         val_acc, test_acc = run(
             args, graph, labels, pred, train_idx, val_idx, test_idx, evaluator
         )
diff --git a/examples/pytorch/ogb/ogbn-proteins/main_proteins_full_dgl.py b/examples/pytorch/ogb/ogbn-proteins/main_proteins_full_dgl.py
@@ -168,7 +168,7 @@ def main(args):
         if num_patient_epochs == args["patience"]:
             break
 
-    model.load_state_dict(torch.load(model_path))
+    model.load_state_dict(torch.load(model_path, weights_only=False))
     train_score, val_score, test_score = run_an_eval_epoch(
         graph, splitted_idx, model, evaluator
     )
diff --git a/examples/pytorch/ogb_lsc/MAG240M/train.py b/examples/pytorch/ogb_lsc/MAG240M/train.py
@@ -247,7 +247,7 @@ def test(args, dataset, g, feats, paper_offset):
         0.5,
         "paper",
     ).cuda()
-    model.load_state_dict(torch.load(args.model_path))
+    model.load_state_dict(torch.load(args.model_path, weights_only=False))
 
     model.eval()
     correct = total = 0
diff --git a/examples/pytorch/ogb_lsc/MAG240M/train_multi_gpus.py b/examples/pytorch/ogb_lsc/MAG240M/train_multi_gpus.py
@@ -304,7 +304,7 @@ def test(args, dataset, g, feats, paper_offset):
     ).cuda()
 
     # load ddp's model parameters, we need to remove the name of 'module.'
-    state_dict = torch.load(args.model_path)
+    state_dict = torch.load(args.model_path, weights_only=False)
     new_state_dict = OrderedDict()
     for k, v in state_dict.items():
         name = k[7:]
diff --git a/examples/pytorch/ogb_lsc/PCQM4M/test_inference.py b/examples/pytorch/ogb_lsc/PCQM4M/test_inference.py
@@ -206,7 +206,7 @@ def main():
         raise RuntimeError(f"Checkpoint file not found at {checkpoint_path}")
 
     ## reading in checkpoint
-    checkpoint = torch.load(checkpoint_path)
+    checkpoint = torch.load(checkpoint_path, weights_only=False)
     model.load_state_dict(checkpoint["model_state_dict"])
 
     print("Predicting on test data...")
diff --git a/examples/pytorch/pointcloud/bipointnet/train_cls.py b/examples/pytorch/pointcloud/bipointnet/train_cls.py
@@ -136,7 +136,9 @@ def evaluate(net, test_loader, dev):
 
 net = net.to(dev)
 if args.load_model_path:
-    net.load_state_dict(torch.load(args.load_model_path, map_location=dev))
+    net.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
 
 opt = optim.Adam(net.parameters(), lr=1e-3, weight_decay=1e-4)
 
diff --git a/examples/pytorch/pointcloud/edgeconv/main.py b/examples/pytorch/pointcloud/edgeconv/main.py
@@ -115,7 +115,9 @@ def evaluate(model, test_loader, dev):
 model = Model(20, [64, 64, 128, 256], [512, 512, 256], 40)
 model = model.to(dev)
 if args.load_model_path:
-    model.load_state_dict(torch.load(args.load_model_path, map_location=dev))
+    model.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
 
 opt = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)
 
diff --git a/examples/pytorch/pointcloud/pct/train_cls.py b/examples/pytorch/pointcloud/pct/train_cls.py
diff --git a/examples/pytorch/pointcloud/pct/train_partseg.py b/examples/pytorch/pointcloud/pct/train_partseg.py
diff --git a/examples/pytorch/pointcloud/point_transformer/train_cls.py b/examples/pytorch/pointcloud/point_transformer/train_cls.py
diff --git a/examples/pytorch/pointcloud/point_transformer/train_partseg.py b/examples/pytorch/pointcloud/point_transformer/train_partseg.py
diff --git a/examples/pytorch/pointcloud/pointnet/train_cls.py b/examples/pytorch/pointcloud/pointnet/train_cls.py
diff --git a/examples/pytorch/pointcloud/pointnet/train_partseg.py b/examples/pytorch/pointcloud/pointnet/train_partseg.py
diff --git a/examples/pytorch/rgcn/link.py b/examples/pytorch/rgcn/link.py
diff --git a/examples/pytorch/rrn/sudoku_solver.py b/examples/pytorch/rrn/sudoku_solver.py
diff --git a/examples/pytorch/rrn/train_sudoku.py b/examples/pytorch/rrn/train_sudoku.py
diff --git a/examples/pytorch/stgcn_wave/main.py b/examples/pytorch/stgcn_wave/main.py
diff --git a/python/dgl/distributed/partition.py b/python/dgl/distributed/partition.py
diff --git a/python/dgl/graphbolt/impl/ondisk_dataset.py b/python/dgl/graphbolt/impl/ondisk_dataset.py
diff --git a/python/dgl/graphbolt/impl/torch_based_feature_store.py b/python/dgl/graphbolt/impl/torch_based_feature_store.py
diff --git a/python/dgl/graphbolt/internal/utils.py b/python/dgl/graphbolt/internal/utils.py
diff --git a/tests/python/pytorch/graphbolt/impl/test_fused_csc_sampling_graph.py b/tests/python/pytorch/graphbolt/impl/test_fused_csc_sampling_graph.py
diff --git a/tests/python/pytorch/graphbolt/impl/test_ondisk_dataset.py b/tests/python/pytorch/graphbolt/impl/test_ondisk_dataset.py
diff --git a/tests/python/pytorch/graphbolt/internal/test_utils.py b/tests/python/pytorch/graphbolt/internal/test_utils.py
