Skip to content
Snippets Groups Projects
Commit a6182eaa authored by Leyuan Wang's avatar Leyuan Wang Committed by Tianqi Chen
Browse files

unroll condition changed for cuda (#232)

parent a0b0546a
No related branches found
No related tags found
No related merge requests found
......@@ -49,8 +49,12 @@ def main():
else:
raise ValueError('no benchmark prepared for {}.'.format(args.model))
if args.target == "cuda":
unroll = 1400
else:
unroll = 128
with nnvm.compiler.build_config(opt_level=opt_level):
with tvm.build_config(auto_unroll_max_step=128,
with tvm.build_config(auto_unroll_max_step=unroll,
unroll_explicit=(args.target != "cuda")):
graph, lib, params = nnvm.compiler.build(
net, args.target, shape={"data": data_shape}, params=params)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment