[Example][Bugfix] Bugfix for dgi example (dmlc#4201)

zephyr0703 · Jul 2, 2022 · 0f0e7c7 · 0f0e7c7
1 parent dcf1699
commit 0f0e7c7
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/examples/pytorch/dgi/train.py b/examples/pytorch/dgi/train.py
@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import dgl
 from dgl import DGLGraph
 from dgl.data import register_data_args, load_data
 from dgi import DGI, Classifier
@@ -49,8 +50,8 @@ def main(args):
 
     # add self loop
     if args.self_loop:
-        g.remove_edges_from(nx.selfloop_edges(g))
-        g.add_edges_from(zip(g.nodes(), g.nodes()))
+        g = dgl.remove_self_loop(g)
+        g = dgl.add_self_loop(g)
     n_edges = g.number_of_edges()
 
     if args.gpu >= 0: