Commit a6182eaa by Leyuan Wang Committed by Tianqi Chen

unroll condition changed for cuda (#232)

parent a0b0546a
......@@ -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)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment