From 17a5f882c8022cdbe5101d514692143a083cf444 Mon Sep 17 00:00:00 2001 From: "padarn.wilson" Date: Sat, 9 Jul 2022 17:28:33 +0800 Subject: [PATCH] add ast for lists --- benchmark/loader/neighbor_loader.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/benchmark/loader/neighbor_loader.py b/benchmark/loader/neighbor_loader.py index 7460b743a0c72..1bc1d874062df 100644 --- 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)