add ast for lists

Padarn · Jul 9, 2022 · 17a5f88 · 17a5f88
1 parent 4f0df80
commit 17a5f88
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/benchmark/loader/neighbor_loader.py b/benchmark/loader/neighbor_loader.py
@@ -1,4 +1,5 @@
 import argparse
+import ast
 import os.path as osp
 from timeit import default_timer
 
@@ -88,9 +89,9 @@ def run(args: argparse.ArgumentParser) -> None:
     add('--eval-batch-sizes', default=[16384, 8192, 4096, 2048, 1024, 512],
         type=int, nargs='+')
     add('--homo-neighbor_sizes', default=[[10, 5], [15, 10, 5], [20, 15, 10]],
-        type=int, nargs='+')
-    add('--hetero-neighbor_sizes', default=[[5], [10], [10, 5]], type=int,
-        nargs='+')
+        type=ast.literal_eval)
+    add('--hetero-neighbor_sizes', default=[[5], [10], [10, 5]],
+        type=ast.literal_eval)
     add('--num-workers', default=0)
     add('--runs', default=3)