Update eval_ogbn_mag.py

acbull · Dec 8, 2020 · fa0ec29 · fa0ec29
1 parent d40b377
commit fa0ec29
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/ogbn-mag/eval_ogbn_mag.py b/ogbn-mag/eval_ogbn_mag.py
@@ -34,7 +34,7 @@
                     help='Number of process to sample subgraph')  
 parser.add_argument('--n_batch', type=int, default=32,
                     help='Number of batch (sampled graphs) for each epoch') 
-parser.add_argument('--batch_size', type=int, default=256,
+parser.add_argument('--batch_size', type=int, default=128,
                     help='Number of output nodes for training')   
 
 
@@ -48,6 +48,8 @@
                     help='Number of attention head')
 parser.add_argument('--n_layers', type=int, default=4,
                     help='Number of GNN layers')
+parser.add_argument('--cuda', type=int, default=2,
+                    help='cuda')
 parser.add_argument('--dropout', type=float, default=0.2,
                     help='Dropout ratio')
 parser.add_argument('--sample_depth', type=int, default=6,