Skip to content

T
tic
Commit 73775f45 by Yizhi Liu Committed by Tianqi Chen
Options

AlterOpLayout with tvm.target (#463) 

* AlterOpLayout with tvm.target

* fix test
parent 343eb82c
Showing with 2 additions and 1 deletions
nnvm/python/nnvm/compiler/build_module.py
...@@ -264,7 +264,8 @@ def build(graph, target=None, shape=None, dtype="float32", ...@@ -264,7 +264,8 @@ def build(graph, target=None, shape=None, dtype="float32",
if _all_var_init: if _all_var_init:
init_var = initialize_variables(shape, dtype) init_var = initialize_variables(shape, dtype)
# Apply optimization # Apply optimization
graph = optimize(graph, shape, dtype, layout) with target:
graph = optimize(graph, shape, dtype, layout)
# Precompute prune # Precompute prune
if params and cfg.pass_enabled("PrecomputePrune"): if params and cfg.pass_enabled("PrecomputePrune"):
graph, params = precompute_prune(graph, params) graph, params = precompute_prune(graph, 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