AlterOpLayout with tvm.target (#463)

* AlterOpLayout with tvm.target * fix test

AlterOpLayout with tvm.target (#463)
* AlterOpLayout with tvm.target * fix test
73775f45 · Yizhi Liu · Tianqi Chen · 343eb82c · 73775f45
Commit 73775f45 authored May 03, 2018 by Yizhi Liu Committed by Tianqi Chen May 29, 2018
Show whitespace changes
Inline Side-by-side

Showing with 1 additions and 0 deletions

nnvm/python/nnvm/compiler/build_module.py
+1 -0

No files found.
--- a/nnvm/python/nnvm/compiler/build_module.py
+++ b/nnvm/python/nnvm/compiler/build_module.py
@@ -264,6 +264,7 @@ def build(graph, target=None, shape=None, dtype="float32",
    if _all_var_init:
        init_var = initialize_variables(shape, dtype)
    # Apply optimization
+    with target:
        graph = optimize(graph, shape, dtype, layout)
    # Precompute prune
    if params and cfg.pass_enabled("PrecomputePrune"):