Skip to content

GitLab

Commit 6a0c21bd authored by rusty1s's avatar rusty1s
Browse files

add comment

parent 7e84bf5f
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
README.md
View file @ 6a0c21bd
...@@ -22,9 +22,12 @@ from torch_geometric_autoscale import metis, permute, SubgraphLoader ...@@ -22,9 +22,12 @@ from torch_geometric_autoscale import metis, permute, SubgraphLoader
class GNN(ScalableGNN): class GNN(ScalableGNN):
def __init__(self, num_nodes, in_channels, hidden_channels, out_channels, num_layers): def __init__(self, num_nodes, in_channels, hidden_channels, out_channels, num_layers):
# pool_size determines the number of pinned CPU buffers
# buffer_size determines the size of pinned CPU buffers
super(GNN, self).__init__(num_nodes, hidden_channels, num_layers, super(GNN, self).__init__(num_nodes, hidden_channels, num_layers,
pool_size=2, buffer_size=5000) pool_size=2, buffer_size=5000)
self.convs = ModuleList() self.convs = ModuleList()
self.convs.append(SAGEConv(in_channels, hidden_channels)) self.convs.append(SAGEConv(in_channels, hidden_channels))
for _ in range(num_layers - 2): for _ in range(num_layers - 2):
......
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