Skip to content

Commit f1f521f

Browse files
jspark1105facebook-github-bot
authored andcommitted
make bench_gen.py work for 3d conv (pytorch#12433)
Summary: Pull Request resolved: pytorch#12433 To test 3d conv, we need to pass lists in spec argument. We also don't want to set use_cudnn=True which is the default in brew. Reviewed By: llyfacebook, csummersea Differential Revision: D10234315 fbshipit-source-id: 96a39992a97e020d6e9dac103e6d64df0cc1020b
1 parent 00aedfc commit f1f521f

File tree

1 file changed

+4
-6
lines changed

1 file changed

+4
-6
lines changed

binaries/bench_gen/bench_gen.py

+4-6
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from __future__ import unicode_literals
77

88
import argparse
9+
import ast
910

1011
from caffe2.python.model_helper import ModelHelper
1112
from caffe2.python.predictor import mobile_exporter
@@ -15,18 +16,15 @@
1516
def parse_kwarg(kwarg_str):
1617
key, value = kwarg_str.split('=')
1718
try:
18-
value = int(value)
19+
value = ast.literal_eval(value)
1920
except ValueError:
20-
try:
21-
value = float(value)
22-
except ValueError:
23-
pass
21+
pass
2422
return key, value
2523

2624

2725
def main(args):
2826
# User defined keyword arguments
29-
kwargs = {"order": "NCHW"}
27+
kwargs = {"order": "NCHW", "use_cudnn": False}
3028
kwargs.update(dict(args.kwargs))
3129

3230
model = ModelHelper(name=args.benchmark_name)

0 commit comments

Comments
 (0)