Skip to content

GitLab

Unverified Commit 0f0e7c7f authored by Chang Liu's avatar Chang Liu Committed by GitHub
Browse files

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

parent dcf16992
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
examples/pytorch/dgi/train.py
View file @ 0f0e7c7f
...@@ -4,6 +4,7 @@ import networkx as nx ...@@ -4,6 +4,7 @@ import networkx as nx
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
import dgl
from dgl import DGLGraph from dgl import DGLGraph
from dgl.data import register_data_args, load_data from dgl.data import register_data_args, load_data
from dgi import DGI, Classifier from dgi import DGI, Classifier
...@@ -49,8 +50,8 @@ def main(args): ...@@ -49,8 +50,8 @@ def main(args):
# add self loop # add self loop
if args.self_loop: if args.self_loop:
g.remove_edges_from(nx.selfloop_edges(g)) g = dgl.remove_self_loop(g)
g.add_edges_from(zip(g.nodes(), g.nodes())) g = dgl.add_self_loop(g)
n_edges = g.number_of_edges() n_edges = g.number_of_edges()
if args.gpu >= 0: if args.gpu >= 0:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment