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