add num_hidden param.

dmlc · Aug 21, 2024 · 9b7d71d · 9b7d71d
1 parent 77c2bb9
commit 9b7d71d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/examples/graphbolt/pyg/hetero/node_classification.py b/examples/graphbolt/pyg/hetero/node_classification.py
@@ -369,6 +369,7 @@ def parse_args():
         default=0.001,
         help="Learning rate for optimization.",
     )
+    parser.add_argument("--num-hidden", type=int, default=1024)
     parser.add_argument(
         "--batch-size", type=int, default=1024, help="Batch size for training."
     )
@@ -524,7 +525,7 @@ def main():
     )
 
     feat_size = features.size("node", "paper", "feat")[0]
-    hidden_channels = 256
+    hidden_channels = args.num_hidden
 
     # Initialize the entity classification model.
     model = EntityClassify(