Unverified Commit 26cbc3fb by Tianqi Chen Committed by GitHub

[CI] Move gpu docker binary to cuda10 (#4229)

* [CI] Move gpu docker binary to cuda10

* Fix the gcn tutorial
parent 18673112
......@@ -45,7 +45,7 @@
//
ci_lint = "tvmai/ci-lint:v0.51"
ci_gpu = "tvmai/ci-gpu:v0.54"
ci_gpu = "tvmai/ci-gpu:v0.55"
ci_cpu = "tvmai/ci-cpu:v0.54"
ci_i386 = "tvmai/ci-i386:v0.52"
......
......@@ -46,6 +46,7 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
import networkx as nx
from dgl.nn.pytorch import GraphConv
class GCN(nn.Module):
......@@ -88,7 +89,7 @@ def load_dataset(dataset="cora"):
# Remove self-loops to avoid duplicate passing of a node's feature to itself
g = data.graph
g.remove_edges_from(g.selfloop_edges())
g.remove_edges_from(nx.selfloop_edges(g))
g.add_edges_from(zip(g.nodes, g.nodes))
return g, data
......@@ -110,7 +111,7 @@ def evaluate(data, logits):
Parameters
----------
dataset: str
Name of dataset. You can choose from ['cora', 'citeseer', 'pubmed'].
Name of dataset. You can choose from ['cora', 'citeseer', 'pubmed'].
num_layer: int
number of hidden layers
......@@ -251,7 +252,7 @@ def GraphConv(layer_name,
######################################################################
# Prepare the parameters needed in the GraphConv layers
# ------------------
#
#
import numpy as np
import networkx as nx
......
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